nodes/thomaskippster-comfymodeldownloader
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor package structure to de.tki.comfymodels, fix compilation errors, and prioritize imported model list in search

Browse Source
This commit is contained in:
Thomas
2026-04-19 11:30:23 +02:00
parent 19d7eb98e3
commit 9d722bf251
22 changed files with 1715 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+13 -9
View File
@@ -1,8 +1,8 @@
# Build artifacts
target/
.gradle/
*.jar
*.class
dependency-reduced-pom.xml
# IDE files
.idea/
@@ -10,16 +10,20 @@ target/
.settings/
.project
.classpath
*.shfbproj
*.user
*.suo
*.userosscache
*.sln.docstates
# OS files
Thumbs.db
.DS_Store
# Temporary files
temp/
dependency-reduced-pom.xml
# --- Project Specific ---
# Sensitive API keys
app_settings.json
# User Model List (internal DB)
uploaded_models.json
# Local config & assets
workflow.json
*.png
*.bat
!src/main/resources/known_good_urls.json
pom.xml
+24 -2
View File
@@ -4,7 +4,7 @@
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<groupId>com.example</groupId>
<groupId>de.tki</groupId>
<artifactId>comfymodeldownloader</artifactId>
<version>1.0-SNAPSHOT</version>
@@ -20,6 +20,28 @@
<artifactId>json</artifactId>
<version>20240303</version>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-context</artifactId>
<version>6.1.5</version>
</dependency>
<dependency>
<groupId>jakarta.annotation</groupId>
<artifactId>jakarta.annotation-api</artifactId>
<version>2.1.1</version>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter-api</artifactId>
<version>5.10.2</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter-engine</artifactId>
<version>5.10.2</version>
<scope>test</scope>
</dependency>
</dependencies>
<build>
@@ -42,7 +64,7 @@
<configuration>
<transformers>
<transformer implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
<mainClass>com.example.comfymodels.Main</mainClass>
<mainClass>de.tki.comfymodels.Main</mainClass>
</transformer>
</transformers>
</configuration>
src/main/java/com/example/comfymodels/Main.java
-499
View File
@@ -1,499 +0,0 @@
package com.example.comfymodels;
import org.json.JSONArray;
import org.json.JSONObject;
import javax.swing.*;
import javax.swing.filechooser.FileNameExtensionFilter;
import javax.swing.table.DefaultTableModel;
import java.awt.*;
import java.io.*;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.file.*;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
public class Main extends JFrame {
private JTextField modelsPathField;
private JTextArea jsonInputArea;
private JTable modelTable;
private DefaultTableModel tableModel;
private JProgressBar progressBar;
private JLabel statusLabel;
private JButton downloadButton;
private JButton pauseButton;
private JButton stopButton;
private List<ModelInfo> modelsToDownload = new ArrayList<>();
private final Set<String> processedModelNames = new HashSet<>();
private final ExecutorService executor = Executors.newSingleThreadExecutor();
private final HttpClient httpClient = HttpClient.newBuilder().followRedirects(HttpClient.Redirect.ALWAYS).build();
private volatile boolean isPaused = false;
private volatile boolean isStopped = false;
private static final Pattern MODEL_FILE_PATTERN = Pattern.compile("([^\\\\/]+\\.(?:safetensors|ckpt|pt|bin|pth|onnx))", Pattern.CASE_INSENSITIVE);
public Main() {
setTitle("ComfyUI Model Download Manager (Pause/Resume/Stop)");
setSize(1200, 850);
setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
setLayout(new BorderLayout());
// Top Panel: Models Directory
JPanel topPanel = new JPanel(new BorderLayout());
topPanel.setBorder(BorderFactory.createTitledBorder("Models Directory"));
modelsPathField = new JTextField("C:\\pinokio\\api\\comfy.git\\app\\models");
JButton browsePathButton = new JButton("Browse...");
browsePathButton.addActionListener(e -> chooseDirectory());
topPanel.add(modelsPathField, BorderLayout.CENTER);
topPanel.add(browsePathButton, BorderLayout.EAST);
// Center Panel: JSON Input and Model List
JSplitPane splitPane = new JSplitPane(JSplitPane.VERTICAL_SPLIT);
// JSON Input
JPanel jsonPanel = new JPanel(new BorderLayout());
jsonPanel.setBorder(BorderFactory.createTitledBorder("ComfyUI JSON Workflow"));
jsonInputArea = new JTextArea();
JScrollPane jsonScroll = new JScrollPane(jsonInputArea);
JPanel jsonButtons = new JPanel(new FlowLayout(FlowLayout.RIGHT));
JButton loadJsonButton = new JButton("Load File...");
loadJsonButton.addActionListener(e -> chooseJsonFile());
JButton analyzeButton = new JButton("Analyze Models");
analyzeButton.addActionListener(e -> analyzeJson());
jsonButtons.add(loadJsonButton);
jsonButtons.add(analyzeButton);
jsonPanel.add(jsonScroll, BorderLayout.CENTER);
jsonPanel.add(jsonButtons, BorderLayout.SOUTH);
// Model List
String[] columnNames = {"Select", "Type", "Name", "URL", "Status"};
tableModel = new DefaultTableModel(columnNames, 0) {
@Override
public Class<?> getColumnClass(int columnIndex) {
return columnIndex == 0 ? Boolean.class : String.class;
}
@Override
public boolean isCellEditable(int row, int column) {
return column == 0;
}
};
modelTable = new JTable(tableModel);
modelTable.getColumnModel().getColumn(0).setMaxWidth(60);
modelTable.getColumnModel().getColumn(1).setPreferredWidth(100);
modelTable.getColumnModel().getColumn(2).setPreferredWidth(250);
modelTable.getColumnModel().getColumn(4).setPreferredWidth(300);
JScrollPane tableScroll = new JScrollPane(modelTable);
JPanel tablePanel = new JPanel(new BorderLayout());
tablePanel.setBorder(BorderFactory.createTitledBorder("Identified Models"));
JPanel selectionPanel = new JPanel(new FlowLayout(FlowLayout.LEFT));
JButton selectAllBtn = new JButton("Select All");
selectAllBtn.addActionListener(e -> setAllSelected(true));
JButton deselectAllBtn = new JButton("Deselect All");
deselectAllBtn.addActionListener(e -> setAllSelected(false));
selectionPanel.add(selectAllBtn);
selectionPanel.add(deselectAllBtn);
tablePanel.add(selectionPanel, BorderLayout.NORTH);
tablePanel.add(tableScroll, BorderLayout.CENTER);
splitPane.setTopComponent(jsonPanel);
splitPane.setBottomComponent(tablePanel);
splitPane.setDividerLocation(300);
// Bottom Panel: Actions and Progress
JPanel bottomPanel = new JPanel(new BorderLayout());
progressBar = new JProgressBar(0, 100);
statusLabel = new JLabel("Ready");
JPanel actionButtons = new JPanel(new FlowLayout(FlowLayout.RIGHT));
downloadButton = new JButton("Start Queue");
downloadButton.setEnabled(false);
downloadButton.addActionListener(e -> startDownloads());
pauseButton = new JButton("Pause");
pauseButton.setEnabled(false);
pauseButton.addActionListener(e -> togglePause());
stopButton = new JButton("Stop");
stopButton.setEnabled(false);
stopButton.addActionListener(e -> stopDownloads());
actionButtons.add(downloadButton);
actionButtons.add(pauseButton);
actionButtons.add(stopButton);
JPanel progressPanel = new JPanel(new GridLayout(2, 1));
progressPanel.add(statusLabel);
progressPanel.add(progressBar);
bottomPanel.add(progressPanel, BorderLayout.CENTER);
bottomPanel.add(actionButtons, BorderLayout.EAST);
add(topPanel, BorderLayout.NORTH);
add(splitPane, BorderLayout.CENTER);
add(bottomPanel, BorderLayout.SOUTH);
}
private void togglePause() {
isPaused = !isPaused;
pauseButton.setText(isPaused ? "Resume" : "Pause");
if (!isPaused) {
statusLabel.setText("Resuming...");
} else {
statusLabel.setText("Paused.");
}
}
private void stopDownloads() {
isStopped = true;
isPaused = false;
pauseButton.setText("Pause");
statusLabel.setText("Stopping queue...");
}
private void setAllSelected(boolean selected) {
for (int i = 0; i < tableModel.getRowCount(); i++) {
tableModel.setValueAt(selected, i, 0);
}
}
private void chooseDirectory() {
JFileChooser chooser = new JFileChooser();
chooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
if (chooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
modelsPathField.setText(chooser.getSelectedFile().getAbsolutePath());
}
}
private void chooseJsonFile() {
JFileChooser chooser = new JFileChooser();
chooser.setFileFilter(new FileNameExtensionFilter("JSON Files", "json"));
if (chooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
try {
String content = Files.readString(chooser.getSelectedFile().toPath());
jsonInputArea.setText(content);
analyzeJson();
} catch (IOException e) {
JOptionPane.showMessageDialog(this, "Error reading file: " + e.getMessage());
}
}
}
private void analyzeJson() {
String jsonText = jsonInputArea.getText();
if (jsonText.isEmpty()) return;
try {
modelsToDownload.clear();
processedModelNames.clear();
tableModel.setRowCount(0);
JSONObject jsonObject = new JSONObject(jsonText);
findModelsMetadata(jsonObject);
scanForModelFiles(jsonObject);
for (ModelInfo info : modelsToDownload) {
tableModel.addRow(new Object[]{true, info.type, info.name, info.url, "Pending"});
}
downloadButton.setEnabled(!modelsToDownload.isEmpty());
statusLabel.setText("Found " + modelsToDownload.size() + " models.");
} catch (Exception e) {
JOptionPane.showMessageDialog(this, "Error analyzing JSON: " + e.getMessage());
}
}
private void findModelsMetadata(Object obj) {
if (obj instanceof JSONObject) {
JSONObject jo = (JSONObject) obj;
if (jo.has("models") && jo.get("models") instanceof JSONArray) {
JSONArray models = jo.getJSONArray("models");
for (int i = 0; i < models.length(); i++) {
Object mObj = models.get(i);
if (mObj instanceof JSONObject) {
JSONObject m = (JSONObject) mObj;
String name = m.optString("name");
if (name != null && !name.isEmpty()) {
String url = m.optString("url", "MISSING");
String dir = m.optString("directory", "checkpoints");
addModelInfo(new ModelInfo(dir, name, url));
}
}
}
}
for (String key : jo.keySet()) {
findModelsMetadata(jo.get(key));
}
} else if (obj instanceof JSONArray) {
JSONArray ja = (JSONArray) obj;
for (int i = 0; i < ja.length(); i++) {
findModelsMetadata(ja.get(i));
}
}
}
private void scanForModelFiles(Object obj) {
if (obj instanceof JSONObject) {
JSONObject jo = (JSONObject) obj;
for (String key : jo.keySet()) {
Object val = jo.get(key);
if (val instanceof String) {
checkAndAddModelFromFileName((String) val, inferTypeFromKey(key));
} else {
scanForModelFiles(val);
}
}
} else if (obj instanceof JSONArray) {
JSONArray ja = (JSONArray) obj;
for (int i = 0; i < ja.length(); i++) {
Object val = ja.get(i);
if (val instanceof String) {
checkAndAddModelFromFileName((String) val, "checkpoints");
} else {
scanForModelFiles(val);
}
}
}
}
private void checkAndAddModelFromFileName(String text, String type) {
Matcher m = MODEL_FILE_PATTERN.matcher(text);
while (m.find()) {
String fileName = m.group(1);
if (!processedModelNames.contains(fileName)) {
addModelInfo(new ModelInfo(type, fileName, "MISSING (Needs manual URL)"));
}
}
}
private String inferTypeFromKey(String key) {
if (key.contains("ckpt") || key.contains("checkpoint")) return "checkpoints";
if (key.contains("lora")) return "loras";
if (key.contains("vae")) return "vae";
if (key.contains("control") && key.contains("net")) return "controlnet";
if (key.contains("upscale")) return "upscale_models";
if (key.contains("clip")) return "clip";
if (key.contains("unet") || key.contains("diffusion")) return "unet";
return "checkpoints";
}
private void addModelInfo(ModelInfo info) {
for (int i = 0; i < modelsToDownload.size(); i++) {
ModelInfo existing = modelsToDownload.get(i);
if (existing.name.equals(info.name)) {
if (existing.url.startsWith("MISSING") && !info.url.startsWith("MISSING")) {
modelsToDownload.set(i, info);
}
return;
}
}
modelsToDownload.add(info);
processedModelNames.add(info.name);
}
private void startDownloads() {
String baseDir = modelsPathField.getText();
if (baseDir.isEmpty()) {
JOptionPane.showMessageDialog(this, "Please specify a models directory.");
return;
}
isStopped = false;
isPaused = false;
downloadButton.setEnabled(false);
pauseButton.setEnabled(true);
stopButton.setEnabled(true);
executor.submit(() -> {
try {
for (int i = 0; i < tableModel.getRowCount(); i++) {
if (isStopped) {
updateStatus(i, "Stopped");
continue;
}
boolean selected = (Boolean) tableModel.getValueAt(i, 0);
if (!selected) continue;
ModelInfo info = modelsToDownload.get(i);
if (info.url.startsWith("MISSING")) {
updateStatus(i, "Skipped (No URL)");
continue;
}
try {
processDownload(info, baseDir, i);
} catch (Exception e) {
updateStatus(i, "Error: " + e.getMessage());
}
}
} finally {
SwingUtilities.invokeLater(() -> {
statusLabel.setText(isStopped ? "Queue stopped." : "Queue finished.");
downloadButton.setEnabled(true);
pauseButton.setEnabled(false);
stopButton.setEnabled(false);
pauseButton.setText("Pause");
});
}
});
}
private void updateStatus(int row, String status) {
SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, row, 4));
}
private void processDownload(ModelInfo info, String baseDir, int rowIndex) throws Exception {
Path targetDir = Paths.get(baseDir, info.type);
Files.createDirectories(targetDir);
Path targetFile = targetDir.resolve(info.name);
long existingSize = 0;
if (Files.exists(targetFile)) {
existingSize = Files.size(targetFile);
long totalSize = getServerFileSize(info.url);
if (totalSize != -1 && existingSize == totalSize) {
updateStatus(rowIndex, "Already exists (Complete)");
return;
}
final long fTotalSize = totalSize;
final long fExistingSize = existingSize;
int choice = showChoiceDialog(info.name, fExistingSize, fTotalSize);
if (choice == 2) { // Skip
updateStatus(rowIndex, "Skipped");
return;
} else if (choice == 1) { // Overwrite
Files.delete(targetFile);
existingSize = 0;
}
}
downloadWithResume(info, targetFile, existingSize, rowIndex);
}
private long getServerFileSize(String url) {
try {
HttpRequest headRequest = HttpRequest.newBuilder()
.uri(URI.create(url))
.method("HEAD", HttpRequest.BodyPublishers.noBody())
.build();
HttpResponse<Void> response = httpClient.send(headRequest, HttpResponse.BodyHandlers.discarding());
return response.headers().firstValueAsLong("Content-Length").orElse(-1L);
} catch (Exception e) {
return -1;
}
}
private int showChoiceDialog(String name, long current, long total) throws Exception {
final int[] result = new int[1];
String totalStr = total == -1 ? "unknown" : (total / 1024 / 1024) + " MB";
String currentStr = (current / 1024 / 1024) + " MB";
SwingUtilities.invokeAndWait(() -> {
String message = String.format("File '%s' already exists.\nLocal size: %s\nRemote size: %s", name, currentStr, totalStr);
Object[] options = {"Resume", "Overwrite", "Skip"};
result[0] = JOptionPane.showOptionDialog(this, message, "Existing File Found",
JOptionPane.YES_NO_CANCEL_OPTION, JOptionPane.QUESTION_MESSAGE, null, options, options[0]);
});
return result[0];
}
private void downloadWithResume(ModelInfo info, Path targetFile, long startByte, int rowIndex) throws IOException, InterruptedException {
HttpRequest.Builder requestBuilder = HttpRequest.newBuilder().uri(URI.create(info.url));
if (startByte > 0) {
requestBuilder.header("Range", "bytes=" + startByte + "-");
}
HttpRequest request = requestBuilder.build();
HttpResponse<InputStream> response = httpClient.send(request, HttpResponse.BodyHandlers.ofInputStream());
if (response.statusCode() != 200 && response.statusCode() != 206) {
throw new IOException("HTTP " + response.statusCode());
}
long contentLen = response.headers().firstValueAsLong("Content-Length").orElse(-1L);
long totalBytes = startByte + (contentLen != -1 ? contentLen : 0);
try (InputStream is = response.body();
OutputStream os = new FileOutputStream(targetFile.toFile(), startByte > 0)) {
byte[] buffer = new byte[16384];
long downloaded = startByte;
long lastUpdate = System.currentTimeMillis();
long lastDownloaded = startByte;
int read;
while ((read = is.read(buffer)) != -1) {
if (isStopped) {
updateStatus(rowIndex, "Stopped at " + (downloaded * 100 / (totalBytes > 0 ? totalBytes : 1)) + "%");
return;
}
while (isPaused) {
Thread.sleep(500);
if (isStopped) return;
}
os.write(buffer, 0, read);
downloaded += read;
long now = System.currentTimeMillis();
if (now - lastUpdate > 500) {
long diff = downloaded - lastDownloaded;
double speed = diff / ((now - lastUpdate) / 1000.0) / 1024.0 / 1024.0;
final long fDownloaded = downloaded;
final double fSpeed = speed;
final int progress = totalBytes > 0 ? (int) (fDownloaded * 100 / totalBytes) : 0;
SwingUtilities.invokeLater(() -> {
progressBar.setValue(progress);
String statusText = String.format("Downloading: %d%% (%.2f MB/s)", progress, fSpeed);
tableModel.setValueAt(statusText, rowIndex, 4);
statusLabel.setText("Current: " + info.name + " [" + statusText + "]");
});
lastUpdate = now;
lastDownloaded = downloaded;
}
}
}
if (!isStopped) updateStatus(rowIndex, "Finished");
}
public static void main(String[] args) {
SwingUtilities.invokeLater(() -> {
Main main = new Main();
main.setLocationRelativeTo(null);
main.setVisible(true);
});
}
static class ModelInfo {
String type;
String name;
String url;
public ModelInfo(String type, String name, String url) {
this.type = type;
this.name = name;
this.url = url;
}
}
}
src/main/java/de/tki/comfymodels/BatchVerification.java
+88
View File
@@ -0,0 +1,88 @@
package de.tki.comfymodels;
import de.tki.comfymodels.domain.ModelInfo;
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.*;
import java.util.List;
import java.util.zip.InflaterInputStream;
public class BatchVerification {
public static void main(String[] args) throws Exception {
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
Path root = Paths.get("C:/pinokio/api/comfy.git/workflows/ComfyUI_examples");
System.out.println("=== BATCH VERIFICATION START ===");
Files.walk(root)
.filter(p -> p.toString().endsWith(".json") || p.toString().endsWith(".png"))
.limit(30) // Stichprobe von 30 verschiedenen Workflows
.forEach(path -> {
try {
System.out.println("\nFILE: " + path.getFileName());
String content;
if (path.toString().endsWith(".png")) {
content = extractWorkflowFromPng(path.toFile());
} else {
content = Files.readString(path);
}
if (content == null) {
System.out.println(" [!] No workflow found in PNG");
return;
}
List<ModelInfo> results = analyzer.analyze(content, path.getFileName().toString());
if (results.isEmpty()) {
System.out.println(" [?] No models detected");
} else {
for (ModelInfo info : results) {
System.out.printf(" [%s] %s | URL: %s | Source: %s\n",
info.getType(), info.getName(),
info.getUrl().equals("MISSING") ? "Missing" : "FOUND",
info.getPopularity());
}
}
} catch (Exception e) {
System.out.println(" [ERROR] " + e.getMessage());
}
});
System.out.println("\n=== BATCH VERIFICATION END ===");
}
private static String extractWorkflowFromPng(File file) {
try (DataInputStream dis = new DataInputStream(new FileInputStream(file))) {
if (dis.readLong() != 0x89504E470D0A1A0AL) return null;
while (true) {
int length = dis.readInt();
byte[] typeBytes = new byte[4];
dis.readFully(typeBytes);
String type = new String(typeBytes, StandardCharsets.US_ASCII);
if ("tEXt".equals(type) || "iTXt".equals(type) || "zTXt".equals(type)) {
byte[] data = new byte[length];
dis.readFully(data);
int nullPos = 0; while (nullPos < data.length && data[nullPos] != 0) nullPos++;
String key = new String(data, 0, nullPos, StandardCharsets.UTF_8);
String value = null;
if ("tEXt".equals(type)) value = new String(data, nullPos + 1, data.length - (nullPos + 1), StandardCharsets.UTF_8);
else if ("zTXt".equals(type)) {
try (InflaterInputStream iis = new InflaterInputStream(new ByteArrayInputStream(data, nullPos + 2, data.length - (nullPos + 2)))) { value = new String(iis.readAllBytes(), StandardCharsets.UTF_8); }
} else if ("iTXt".equals(type)) {
int currentPos = nullPos + 3; int nullCount = 0;
while (currentPos < data.length && nullCount < 2) { if (data[currentPos] == 0) nullCount++; currentPos++; }
if (currentPos < data.length) {
if (data[nullPos + 1] != 0) {
try (InflaterInputStream iis = new InflaterInputStream(new ByteArrayInputStream(data, currentPos, data.length - currentPos))) { value = new String(iis.readAllBytes(), StandardCharsets.UTF_8); }
} else value = new String(data, currentPos, data.length - currentPos, StandardCharsets.UTF_8);
}
}
if ("workflow".equalsIgnoreCase(key) || "prompt".equalsIgnoreCase(key)) return value;
dis.readInt();
} else if ("IEND".equals(type)) break; else dis.skipBytes(length + 4);
}
} catch (Exception e) {}
return null;
}
}
src/main/java/de/tki/comfymodels/Main.java
+324
View File
@@ -0,0 +1,324 @@
package de.tki.comfymodels;
import de.tki.comfymodels.domain.ModelInfo;
import de.tki.comfymodels.service.IDownloadManager;
import de.tki.comfymodels.service.IModelAnalyzer;
import de.tki.comfymodels.service.IWorkflowService;
import de.tki.comfymodels.service.IModelSearchService;
import de.tki.comfymodels.service.impl.ConfigService;
import de.tki.comfymodels.service.impl.GeminiAIService;
import de.tki.comfymodels.service.impl.ModelListService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.Configuration;
import org.springframework.stereotype.Component;
import javax.swing.*;
import javax.swing.table.DefaultTableModel;
import java.awt.*;
import java.awt.datatransfer.DataFlavor;
import java.awt.dnd.*;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.List;
import java.util.Arrays;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.zip.InflaterInputStream;
@Component
public class Main extends JFrame {
private final IModelAnalyzer analyzer;
private final IDownloadManager downloadManager;
private final IWorkflowService workflowService;
private final IModelSearchService searchService;
@Autowired
private ConfigService configService;
@Autowired
private GeminiAIService geminiService;
@Autowired
private ModelListService modelListService;
private JTextField modelsPathField;
private JPasswordField geminiKeyField;
private JPasswordField hfTokenField;
private JLabel activeAiModelLabel;
private JTextArea jsonInputArea;
private DefaultTableModel tableModel;
private JProgressBar progressBar;
private JLabel statusLabel;
private JButton downloadButton, pauseButton, stopButton;
private List<ModelInfo> modelsToDownload;
private String currentFileName = "input.json";
public Main(IModelAnalyzer analyzer, IDownloadManager downloadManager,
IWorkflowService workflowService, IModelSearchService searchService) {
this.analyzer = analyzer;
this.downloadManager = downloadManager;
this.workflowService = workflowService;
this.searchService = searchService;
}
@Autowired
public void start() {
initUI();
updateAiModelDisplay();
SwingUtilities.invokeLater(() -> setVisible(true));
}
private void updateAiModelDisplay() {
new Thread(() -> {
String model = geminiService.discoverBestModel();
SwingUtilities.invokeLater(() -> activeAiModelLabel.setText("Active AI: " + model));
}).start();
}
private void initUI() {
setTitle("ComfyUI Downloader (ULTRA-VALIDATION)");
setSize(1450, 1000);
setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
setLayout(new BorderLayout());
JPanel settingsPanel = new JPanel(new GridLayout(4, 1));
JPanel pathRow = new JPanel(new BorderLayout());
pathRow.setBorder(BorderFactory.createTitledBorder("Models Directory"));
modelsPathField = new JTextField("C:\\pinokio\\api\\comfy.git\\app\\models");
JButton browsePathBtn = new JButton("Set Folder...");
browsePathBtn.addActionListener(e -> {
JFileChooser chooser = new JFileChooser();
chooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
if (chooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
modelsPathField.setText(chooser.getSelectedFile().getAbsolutePath());
}
});
pathRow.add(modelsPathField, BorderLayout.CENTER);
pathRow.add(browsePathBtn, BorderLayout.EAST);
JPanel apiRow = new JPanel(new BorderLayout());
apiRow.setBorder(BorderFactory.createTitledBorder("Gemini AI API Key"));
geminiKeyField = new JPasswordField(configService.getGeminiApiKey());
JButton saveGeminiBtn = new JButton("Save & Discover");
saveGeminiBtn.addActionListener(e -> {
configService.setGeminiApiKey(new String(geminiKeyField.getPassword()));
updateAiModelDisplay();
});
apiRow.add(geminiKeyField, BorderLayout.CENTER);
apiRow.add(saveGeminiBtn, BorderLayout.EAST);
JPanel hfRow = new JPanel(new BorderLayout());
hfRow.setBorder(BorderFactory.createTitledBorder("Hugging Face Access Token"));
hfTokenField = new JPasswordField(configService.getHfToken());
JButton saveHfBtn = new JButton("Save HF Token");
saveHfBtn.addActionListener(e -> {
configService.setHfToken(new String(hfTokenField.getPassword()));
statusLabel.setText("HF Token saved.");
analyzeJsonContent(); // Re-validate sizes with token
});
hfRow.add(hfTokenField, BorderLayout.CENTER);
hfRow.add(saveHfBtn, BorderLayout.EAST);
JPanel infoRow = new JPanel(new FlowLayout(FlowLayout.LEFT));
activeAiModelLabel = new JLabel("Active AI: Detecting...");
activeAiModelLabel.setFont(new Font("SansSerif", Font.BOLD, 12));
activeAiModelLabel.setForeground(new Color(0, 100, 200));
infoRow.add(activeAiModelLabel);
settingsPanel.add(pathRow);
settingsPanel.add(apiRow);
settingsPanel.add(hfRow);
settingsPanel.add(infoRow);
JSplitPane splitPane = new JSplitPane(JSplitPane.VERTICAL_SPLIT);
JPanel jsonPanel = new JPanel(new BorderLayout());
jsonPanel.setBorder(BorderFactory.createTitledBorder("Workflow (Drag & Drop JSON/PNG)"));
jsonInputArea = new JTextArea();
setupDragAndDrop(jsonInputArea);
JScrollPane jsonScroll = new JScrollPane(jsonInputArea);
JPanel jsonButtons = new JPanel(new FlowLayout(FlowLayout.RIGHT));
JButton loadJsonBtn = new JButton("Load Workflow...");
loadJsonBtn.addActionListener(e -> {
FileDialog fd = new FileDialog(this, "Select Workflow", FileDialog.LOAD);
fd.setVisible(true);
if (fd.getFile() != null) loadFile(new File(fd.getDirectory(), fd.getFile()));
});
JButton importModelListBtn = new JButton("Import Model List (JSON)...");
importModelListBtn.addActionListener(e -> {
FileDialog fd = new FileDialog(this, "Select Model List JSON", FileDialog.LOAD);
fd.setVisible(true);
if (fd.getFile() != null) {
try {
modelListService.importJson(new File(fd.getDirectory(), fd.getFile()));
JOptionPane.showMessageDialog(this, "Model list imported successfully! (" + modelListService.getModels().size() + " models)");
analyzeJsonContent(); // Re-analyze current workflow with new list
} catch (Exception ex) {
JOptionPane.showMessageDialog(this, "Error importing model list: " + ex.getMessage());
}
}
});
JButton analyzeBtn = new JButton("Deep Search (Strict Validation)");
analyzeBtn.addActionListener(e -> {
analyzeJsonContent();
searchMissingOnline();
});
jsonButtons.add(loadJsonBtn);
jsonButtons.add(importModelListBtn);
jsonButtons.add(analyzeBtn);
jsonPanel.add(jsonScroll, BorderLayout.CENTER);
jsonPanel.add(jsonButtons, BorderLayout.SOUTH);
String[] columnNames = {"Select", "Type", "Name", "Size", "AI Trust & Source", "URL", "Status"};
tableModel = new DefaultTableModel(columnNames, 0) {
@Override public Class<?> getColumnClass(int c) { return c == 0 ? Boolean.class : String.class; }
@Override public boolean isCellEditable(int r, int c) { return c == 0; }
};
JTable modelTable = new JTable(tableModel);
modelTable.getColumnModel().getColumn(0).setMaxWidth(60);
modelTable.getColumnModel().getColumn(4).setMinWidth(250);
JScrollPane tableScroll = new JScrollPane(modelTable);
JPanel tablePanel = new JPanel(new BorderLayout());
tablePanel.setBorder(BorderFactory.createTitledBorder("Identified Models"));
tablePanel.add(tableScroll, BorderLayout.CENTER);
splitPane.setTopComponent(jsonPanel);
splitPane.setBottomComponent(tablePanel);
splitPane.setDividerLocation(300);
JPanel bottomPanel = new JPanel(new BorderLayout());
progressBar = new JProgressBar(0, 100);
statusLabel = new JLabel("Ready");
JPanel actionButtons = new JPanel(new FlowLayout(FlowLayout.RIGHT));
downloadButton = new JButton("Start Queue");
downloadButton.setEnabled(false);
downloadButton.addActionListener(e -> {
downloadButton.setEnabled(false);
downloadManager.startQueue(modelsToDownload, modelsPathField.getText(),
(idx, status) -> SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, idx, 6)),
() -> SwingUtilities.invokeLater(() -> {
downloadButton.setEnabled(true);
statusLabel.setText("Queue finished.");
})
);
});
pauseButton = new JButton("Pause");
pauseButton.addActionListener(e -> {
downloadManager.togglePause();
pauseButton.setText(downloadManager.isPaused() ? "Resume" : "Pause");
});
stopButton = new JButton("Stop");
stopButton.addActionListener(e -> downloadManager.stop());
actionButtons.add(downloadButton);
actionButtons.add(pauseButton);
actionButtons.add(stopButton);
JPanel progressPanel = new JPanel(new GridLayout(2, 1));
progressPanel.add(statusLabel);
progressPanel.add(progressBar);
bottomPanel.add(progressPanel, BorderLayout.CENTER);
bottomPanel.add(actionButtons, BorderLayout.EAST);
add(settingsPanel, BorderLayout.NORTH);
add(splitPane, BorderLayout.CENTER);
add(bottomPanel, BorderLayout.SOUTH);
}
private void setupDragAndDrop(JTextArea area) {
new DropTarget(area, new DropTargetListener() {
public void dragEnter(DropTargetDragEvent dtde) {}
public void dragOver(DropTargetDragEvent dtde) {}
public void dropActionChanged(DropTargetDragEvent dtde) {}
public void dragExit(DropTargetEvent dte) {}
public void drop(DropTargetDropEvent dtde) {
try {
dtde.acceptDrop(DnDConstants.ACTION_COPY);
List<File> files = (List<File>) dtde.getTransferable().getTransferData(DataFlavor.javaFileListFlavor);
if (!files.isEmpty()) loadFile(files.get(0));
} catch (Exception e) {}
}
});
}
private void loadFile(File file) {
try {
currentFileName = file.getName();
String content = workflowService.extractWorkflow(file);
jsonInputArea.setText(content);
analyzeJsonContent();
} catch (IOException ex) { JOptionPane.showMessageDialog(this, "Error: " + ex.getMessage()); }
}
private void analyzeJsonContent() {
String text = jsonInputArea.getText();
if (text.isEmpty()) return;
modelsToDownload = analyzer.analyze(text, currentFileName);
tableModel.setRowCount(0);
for (int i = 0; i < modelsToDownload.size(); i++) {
ModelInfo info = modelsToDownload.get(i);
tableModel.addRow(new Object[]{true, info.getType(), info.getName(), info.getSize(), info.getPopularity(), info.getUrl(), info.getUrl().equals("MISSING") ? "Idle" : "✅ Known Good"});
if (!info.getUrl().equals("MISSING")) fetchSizeInBackground(info, i);
}
downloadButton.setEnabled(!modelsToDownload.isEmpty());
}
private void searchMissingOnline() {
if (modelsToDownload == null) return;
int rowCount = tableModel.getRowCount();
boolean[] selected = new boolean[rowCount];
boolean anySelected = false;
for (int i = 0; i < rowCount; i++) {
selected[i] = (Boolean) tableModel.getValueAt(i, 0);
if (selected[i]) anySelected = true;
}
if (!anySelected) {
JOptionPane.showMessageDialog(this, "Please select at least one model in the table to search.");
return;
}
statusLabel.setText("Deep Search in progress...");
searchService.searchOnline(modelsToDownload, selected, jsonInputArea.getText(), currentFileName,
(idx, status) -> SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, idx, 6)),
(idx, info) -> SwingUtilities.invokeLater(() -> {
tableModel.setValueAt(info.getSize(), idx, 3);
tableModel.setValueAt(info.getPopularity(), idx, 4);
tableModel.setValueAt(info.getUrl(), idx, 5);
tableModel.setValueAt("✅ Found", idx, 6);
})
);
}
private void fetchSizeInBackground(ModelInfo info, int rowIndex) {
new Thread(() -> {
long size = searchService.getRemoteSize(info.getUrl());
String s = searchService.formatSize(size);
info.setSize(s);
SwingUtilities.invokeLater(() -> tableModel.setValueAt(s, rowIndex, 3));
}).start();
}
public static void main(String[] args) {
AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(AppConfig.class);
context.getBean(Main.class);
}
@Configuration
@ComponentScan("de.tki.comfymodels")
public static class AppConfig {}
}
src/main/java/de/tki/comfymodels/PngTest.java
+49
View File
@@ -0,0 +1,49 @@
package de.tki.comfymodels;
import java.io.*;
import java.nio.charset.StandardCharsets;
public class PngTest {
public static void main(String[] args) {
File file = new File("flux_canny_model_example.png");
String workflow = extractWorkflowFromPng(file);
if (workflow != null) {
System.out.println("SUCCESS! Extracted workflow length: " + workflow.length());
System.out.println("Snippet: " + workflow.substring(0, Math.min(200, workflow.length())));
} else {
System.out.println("FAILED! No workflow found.");
}
}
private static String extractWorkflowFromPng(File file) {
try (DataInputStream dis = new DataInputStream(new FileInputStream(file))) {
if (dis.readLong() != 0x89504E470D0A1A0AL) return null;
while (true) {
int length = dis.readInt();
byte[] typeBytes = new byte[4];
dis.readFully(typeBytes);
String type = new String(typeBytes, StandardCharsets.US_ASCII);
if ("tEXt".equals(type) || "iTXt".equals(type) || "zTXt".equals(type)) {
byte[] data = new byte[length];
dis.readFully(data);
String content = new String(data, StandardCharsets.UTF_8);
if (content.contains("workflow") || content.contains("prompt")) {
for (int i = 0; i < data.length; i++) {
if (data[i] == 0) {
return new String(data, i + 1, data.length - i - 1, StandardCharsets.UTF_8).trim();
}
}
}
dis.readInt(); // CRC
} else if ("IEND".equals(type)) {
break;
} else {
dis.skipBytes(length + 4);
}
}
} catch (Exception e) {
e.printStackTrace();
}
return null;
}
}
src/main/java/de/tki/comfymodels/domain/ModelInfo.java
+56
View File
@@ -0,0 +1,56 @@
package de.tki.comfymodels.domain;
import java.util.Objects;
public class ModelInfo {
private String type;
private String name;
private String url;
private String size = "Unknown";
private String popularity = ""; // Download-Zahlen zur Sicherheit
private String base;
private String save_path;
private String description;
private String reference;
private String filename;
public ModelInfo(String type, String name, String url) {
this.type = type;
this.name = name;
this.url = url;
}
public String getType() { return type; }
public void setType(String type) { this.type = type; }
public String getName() { return name; }
public String getUrl() { return url; }
public void setUrl(String url) { this.url = url; }
public String getSize() { return size; }
public void setSize(String size) { this.size = size; }
public String getPopularity() { return popularity; }
public void setPopularity(String popularity) { this.popularity = popularity; }
public String getBase() { return base; }
public void setBase(String base) { this.base = base; }
public String getSave_path() { return save_path; }
public void setSave_path(String save_path) { this.save_path = save_path; }
public String getDescription() { return description; }
public void setDescription(String description) { this.description = description; }
public String getReference() { return reference; }
public void setReference(String reference) { this.reference = reference; }
public String getFilename() { return filename; }
public void setFilename(String filename) { this.filename = filename; }
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
ModelInfo info = (ModelInfo) o;
return Objects.equals(name, info.name);
}
@Override
public int hashCode() {
return Objects.hash(name);
}
}
src/main/java/de/tki/comfymodels/service/IDownloadManager.java
+12
View File
@@ -0,0 +1,12 @@
package de.tki.comfymodels.service;
import de.tki.comfymodels.domain.ModelInfo;
import java.util.List;
import java.util.function.BiConsumer;
public interface IDownloadManager {
void startQueue(List<ModelInfo> models, String baseDir, BiConsumer<Integer, String> statusUpdater, Runnable onFinished);
void togglePause();
void stop();
boolean isPaused();
}
src/main/java/de/tki/comfymodels/service/IModelAnalyzer.java
+8
View File
@@ -0,0 +1,8 @@
package de.tki.comfymodels.service;
import de.tki.comfymodels.domain.ModelInfo;
import java.util.List;
public interface IModelAnalyzer {
List<ModelInfo> analyze(String jsonText, String fileName);
}
src/main/java/de/tki/comfymodels/service/IModelSearchService.java
+23
View File
@@ -0,0 +1,23 @@
package de.tki.comfymodels.service;
import de.tki.comfymodels.domain.ModelInfo;
import java.util.function.BiConsumer;
import java.util.List;
public interface IModelSearchService {
/**
* Searches for models online.
* @param modelsToDownload List of models to check.
* @param selectedIndices Boolean array indicating which models are selected for search.
* @param workflowContext The context of the workflow (JSON string).
* @param fileName The name of the file being processed.
* @param onStatusUpdate Callback for status updates (index, status).
* @param onModelFound Callback when a model URL is found (index, modelInfo).
*/
void searchOnline(List<ModelInfo> modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound);
long getRemoteSize(String url);
String formatSize(long bytes);
}
src/main/java/de/tki/comfymodels/service/IWorkflowService.java
+11
View File
@@ -0,0 +1,11 @@
package de.tki.comfymodels.service;
import java.io.File;
import java.io.IOException;
public interface IWorkflowService {
/**
* Extracts the workflow JSON string from a file (JSON or PNG).
*/
String extractWorkflow(File file) throws IOException;
}
src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java
+192
View File
@@ -0,0 +1,192 @@
package de.tki.comfymodels.service.impl;
import de.tki.comfymodels.domain.ModelInfo;
import de.tki.comfymodels.service.IModelAnalyzer;
import org.json.JSONArray;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import jakarta.annotation.PostConstruct;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
@Service
public class ComfyModelAnalyzer implements IModelAnalyzer {
@Autowired(required = false)
private LocalAIService aiService;
@Autowired
private ModelListService modelListService;
@Autowired(required = false)
private GeminiAIService geminiService;
private static final class ExpertPatterns {
static final Pattern FILE = Pattern.compile("([a-zA-Z0-9_\\-\\.\\/]+\\.(?:safetensors|sft|ckpt|pt|bin|pth|onnx|yaml))", Pattern.CASE_INSENSITIVE);
static final Pattern COMP_VAE = Pattern.compile("(vae|tokenizer|autoencoder)", Pattern.CASE_INSENSITIVE);
static final Pattern COMP_CLIP = Pattern.compile("(clip|t5|text_encoder|llama|gemma|mistral)", Pattern.CASE_INSENSITIVE);
static final Pattern COMP_UNET = Pattern.compile("(unet|diffusion|model|transformer|base|transformer|upscale|esrgan)", Pattern.CASE_INSENSITIVE);
}
private final Map<String, String> knownGoodUrls = new HashMap<>();
@PostConstruct
public void init() {
try (InputStream is = getClass().getResourceAsStream("/known_good_urls.json")) {
if (is != null) {
String json = new String(is.readAllBytes(), StandardCharsets.UTF_8);
JSONObject jo = new JSONObject(json);
for (String key : jo.keySet()) {
knownGoodUrls.put(key.toLowerCase(), jo.getString(key));
}
}
} catch (Exception e) {
System.err.println("Failed to load known_good_urls.json: " + e.getMessage());
}
}
@Override
public List<ModelInfo> analyze(String jsonText, String fileName) {
// Fallback falls Spring Injection fehlt (Test-Modus)
if (aiService == null) aiService = new LocalAIService();
if (knownGoodUrls.isEmpty()) init(); // Fallback für Tests ohne Spring Context
List<ModelInfo> results = new ArrayList<>();
if (jsonText == null || jsonText.trim().isEmpty()) return results;
try {
JSONObject jo = parseJson(jsonText);
if (jo == null) return results;
findModelsMetadata(jo, results);
scanForModelFiles(jo, null, results);
applyContext(results, fileName);
for (ModelInfo info : results) {
String low = info.getName().toLowerCase();
Optional<ModelInfo> uploadedMatch = (modelListService != null) ? modelListService.findByFilename(info.getName()) : Optional.empty();
if (uploadedMatch.isPresent()) {
info.setUrl(uploadedMatch.get().getUrl());
info.setPopularity("📂 USER DEFINED LIST");
} else if (knownGoodUrls.containsKey(low)) {
info.setUrl(knownGoodUrls.get(low));
info.setPopularity("🧠 AI VERIFIED OFFICIAL");
} else {
LocalAIService.Prediction prediction = aiService.predictProvider(info.getName());
info.setPopularity(prediction.getLabel());
}
}
} catch (Exception e) {}
return results;
}
private JSONObject parseJson(String text) {
try {
String trimmed = text.trim();
if (trimmed.startsWith("{")) return new JSONObject(trimmed);
int s = trimmed.indexOf("{"), e = trimmed.lastIndexOf("}");
if (s != -1 && e != -1) return new JSONObject(trimmed.substring(s, e + 1));
} catch (Exception e) {}
return null;
}
private void applyContext(List<ModelInfo> res, String fn) {
if (fn == null) return;
String ctx = fn.toLowerCase().split("[_\\- ]")[0];
if (ctx.length() < 3) return;
for (ModelInfo i : res) if (i.getName().toLowerCase().contains(ctx) && i.getPopularity().contains("Community")) i.setPopularity("📌 Context Match: " + ctx);
}
private void findModelsMetadata(Object obj, List<ModelInfo> res) {
if (obj instanceof JSONObject) {
JSONObject jo = (JSONObject) obj;
if (jo.has("models") && jo.get("models") instanceof JSONArray) {
JSONArray ms = jo.getJSONArray("models");
for (int i = 0; i < ms.length(); i++) {
JSONObject m = ms.optJSONObject(i);
if (m != null) {
String n = m.optString("name", m.optString("filename"));
if (n != null && !n.isEmpty()) addModelInfo(res, new ModelInfo(inferTypeFromKey(m.optString("type", "checkpoints")), n, m.optString("url", "MISSING")));
}
}
}
for (String k : jo.keySet()) findModelsMetadata(jo.get(k), res);
} else if (obj instanceof JSONArray) {
JSONArray ja = (JSONArray) obj;
for (int i = 0; i < ja.length(); i++) findModelsMetadata(ja.get(i), res);
}
}
private void scanForModelFiles(Object obj, String ctx, List<ModelInfo> res) {
if (obj instanceof JSONObject) {
JSONObject jo = (JSONObject) obj;
String node = jo.optString("class_type", jo.optString("type", null));
String curCtx = node != null ? (inferTypeFromNode(node) != null ? inferTypeFromNode(node) : ctx) : ctx;
if (jo.has("widgets_values")) {
JSONArray ws = jo.optJSONArray("widgets_values");
if (ws != null) for (int i = 0; i < ws.length(); i++) if (ws.get(i) instanceof String) processVal((String) ws.get(i), curCtx, node, res);
}
for (String k : jo.keySet()) {
if (k.equals("widgets_values")) continue;
Object v = jo.get(k);
if (v instanceof String) processVal((String) v, curCtx, k, res);
else if (v instanceof JSONObject || v instanceof JSONArray) scanForModelFiles(v, curCtx, res);
}
} else if (obj instanceof JSONArray) {
JSONArray ja = (JSONArray) obj;
for (int i = 0; i < ja.length(); i++) scanForModelFiles(ja.get(i), ctx, res);
}
}
private void processVal(String v, String ctx, String hint, List<ModelInfo> res) {
Matcher m = ExpertPatterns.FILE.matcher(v);
while (m.find()) {
String f = m.group(1);
if (f.length() > 3) {
String type = hint != null ? inferTypeFromKey(hint) : null;
if (type == null) type = ctx;
String fLow = f.toLowerCase();
if (ExpertPatterns.COMP_VAE.matcher(fLow).find()) type = "vae";
else if (ExpertPatterns.COMP_CLIP.matcher(fLow).find()) type = "clip";
else if (ExpertPatterns.COMP_UNET.matcher(fLow).find() && (type == null || type.equals("checkpoints"))) type = "unet";
addModelInfo(res, new ModelInfo(type != null ? type : "checkpoints", f, "MISSING"));
}
}
}
private String inferTypeFromNode(String nt) {
nt = nt.toLowerCase();
if (nt.contains("lora")) return "loras";
if (nt.contains("vae")) return "vae";
if (nt.contains("controlnet") || nt.contains("t2i_adapter")) return "controlnet";
if (nt.contains("upscale") || nt.contains("esrgan")) return "upscale_models";
if (nt.contains("clip") || nt.contains("text_encoder") || nt.contains("cliploader") || nt.contains("llama") || nt.contains("mistral")) return "clip";
if (nt.contains("unet") || nt.contains("diffusion") || nt.contains("ltxv") || nt.contains("cosmos") || nt.contains("hunyuan") || nt.contains("wan") || nt.contains("hidream") || nt.contains("aura") || nt.contains("mochi") || nt.contains("svd")) return "unet";
return null;
}
private String inferTypeFromKey(String k) {
k = k.toLowerCase();
if (k.contains("lora")) return "loras";
if (k.contains("vae")) return "vae";
if (k.contains("clip") || k.contains("text_encoder") || k.contains("llama") || k.contains("mistral")) return "clip";
if (k.contains("unet") || k.contains("diffusion") || k.contains("hidream") || k.contains("aura") || k.contains("mochi") || k.contains("svd")) return "unet";
if (k.contains("control") && k.contains("net")) return "controlnet";
if (k.contains("upscale") || k.contains("esrgan")) return "upscale_models";
return null;
}
private void addModelInfo(List<ModelInfo> res, ModelInfo info) {
for (ModelInfo e : res) {
if (e.getName().equalsIgnoreCase(info.getName())) {
if (e.getUrl().equals("MISSING") && !info.getUrl().equals("MISSING")) e.setUrl(info.getUrl());
if (e.getType().equals("checkpoints") && !info.getType().equals("checkpoints")) e.setType(info.getType());
return;
}
}
res.add(info);
}
}
src/main/java/de/tki/comfymodels/service/impl/ConfigService.java
+45
View File
@@ -0,0 +1,45 @@
package de.tki.comfymodels.service.impl;
import org.json.JSONObject;
import org.springframework.stereotype.Service;
import java.io.File;
import java.nio.file.Files;
import java.nio.file.Paths;
@Service
public class ConfigService {
private static final String CONFIG_FILE = "app_settings.json";
private JSONObject settings;
public ConfigService() {
load();
}
private void load() {
try {
File file = new File(CONFIG_FILE);
if (file.exists()) {
settings = new JSONObject(Files.readString(file.toPath()));
} else {
settings = new JSONObject();
}
} catch (Exception e) {
settings = new JSONObject();
}
}
public void save() {
try {
Files.writeString(Paths.get(CONFIG_FILE), settings.toString(4));
} catch (Exception e) {
e.printStackTrace();
}
}
public String getGeminiApiKey() { return settings.optString("gemini_api_key", ""); }
public void setGeminiApiKey(String key) { settings.put("gemini_api_key", key); save(); }
public String getHfToken() { return settings.optString("hf_token", ""); }
public void setHfToken(String token) { settings.put("hf_token", token); save(); }
}
src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java
+121
View File
@@ -0,0 +1,121 @@
package de.tki.comfymodels.service.impl;
import de.tki.comfymodels.domain.ModelInfo;
import de.tki.comfymodels.service.IDownloadManager;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import java.io.*;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.file.*;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.BiConsumer;
@Service
public class DefaultDownloadManager implements IDownloadManager {
private final ExecutorService executor = Executors.newFixedThreadPool(3);
private final HttpClient httpClient = HttpClient.newBuilder().followRedirects(HttpClient.Redirect.ALWAYS).build();
private volatile boolean isPaused = false;
private volatile boolean isStopped = false;
@Autowired
private ConfigService configService;
@Override
public void startQueue(List<ModelInfo> models, String baseDir, BiConsumer<Integer, String> statusUpdater, Runnable onFinished) {
isStopped = false;
isPaused = false;
new Thread(() -> {
try {
java.util.concurrent.CompletableFuture<?>[] futures = new java.util.concurrent.CompletableFuture[models.size()];
for (int i = 0; i < models.size(); i++) {
final int index = i;
ModelInfo info = models.get(i);
futures[i] = java.util.concurrent.CompletableFuture.runAsync(() -> {
if (isStopped) { statusUpdater.accept(index, "Stopped"); return; }
if (info.getUrl() == null || info.getUrl().startsWith("MISSING")) { statusUpdater.accept(index, "Skipped (No URL)"); return; }
try {
Path targetDir = Paths.get(baseDir, info.getType());
Files.createDirectories(targetDir);
downloadWithResume(info, targetDir.resolve(info.getName()), index, statusUpdater);
} catch (Exception e) { statusUpdater.accept(index, "Error: " + e.getMessage()); }
}, executor);
}
java.util.concurrent.CompletableFuture.allOf(futures).join();
} finally { onFinished.run(); }
}).start();
}
private void downloadWithResume(ModelInfo info, Path targetFile, int index, BiConsumer<Integer, String> statusUpdater) throws Exception {
File file = targetFile.toFile();
// ANTI-STUB LOGIK: Wenn Datei < 10KB, löschen wir sie und fangen neu an (da es wahrscheinlich ein LFS-Pointer ist)
if (file.exists() && file.length() < 10240) {
file.delete();
}
long existingFileSize = file.exists() ? file.length() : 0;
HttpRequest.Builder builder = HttpRequest.newBuilder().uri(URI.create(info.getUrl())).header("User-Agent", "Mozilla/5.0");
String hfToken = configService.getHfToken();
if (info.getUrl().contains("huggingface.co") && !hfToken.isEmpty()) {
builder.header("Authorization", "Bearer " + hfToken);
}
if (existingFileSize > 0) builder.header("Range", "bytes=" + existingFileSize + "-");
HttpResponse<InputStream> response = httpClient.send(builder.build(), HttpResponse.BodyHandlers.ofInputStream());
int statusCode = response.statusCode();
if (statusCode == 401 || statusCode == 403) {
statusUpdater.accept(index, "❌ Auth Required (Token?)");
return;
}
long contentLen = response.headers().firstValueAsLong("Content-Length").orElse(0L);
long totalBytes = (statusCode == 206) ? contentLen + existingFileSize : contentLen;
// VERIFIZIERUNG: Falls die Gesamtgröße extrem klein ist, ist es ein LFS Stub
if (totalBytes > 0 && totalBytes < 5000 && info.getName().endsWith(".safetensors")) {
statusUpdater.accept(index, "❌ LFS Stub detected (Token?)");
return;
}
if (statusCode == 200) existingFileSize = 0;
try (InputStream is = response.body(); RandomAccessFile raf = new RandomAccessFile(file, "rw")) {
raf.seek(existingFileSize);
byte[] buffer = new byte[65536];
long downloaded = existingFileSize;
int read;
long lastUpdate = 0;
while ((read = is.read(buffer)) != -1) {
if (isStopped) return;
while (isPaused) { statusUpdater.accept(index, "Paused..."); Thread.sleep(1000); }
raf.write(buffer, 0, read);
downloaded += read;
long now = System.currentTimeMillis();
if (now - lastUpdate > 800) {
statusUpdater.accept(index, "Downloading: " + (totalBytes > 0 ? (downloaded * 100 / totalBytes) : "?") + "% (" + formatSize(downloaded) + ")");
lastUpdate = now;
}
}
}
statusUpdater.accept(index, "Finished");
}
private String formatSize(long bytes) {
if (bytes < 1024 * 1024) return (bytes / 1024) + " KB";
if (bytes < 1024 * 1024 * 1024) return String.format("%.1f MB", bytes / (1024.0 * 1024.0));
return String.format("%.2f GB", bytes / (1024.0 * 1024.0 * 1024.0));
}
@Override public void togglePause() { isPaused = !isPaused; }
@Override public void stop() { isStopped = true; isPaused = false; }
@Override public boolean isPaused() { return isPaused; }
}
src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java
+120
View File
@@ -0,0 +1,120 @@
package de.tki.comfymodels.service.impl;
import org.json.JSONArray;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
@Service
public class GeminiAIService {
@Autowired
private ConfigService configService;
private final HttpClient client = HttpClient.newBuilder()
.connectTimeout(Duration.ofSeconds(15))
.build();
private String activeModel = "gemini-1.5-flash";
public String getActiveModel() {
return activeModel;
}
private static final List<String> MODEL_PRIORITY = Arrays.asList(
"gemini-3.1-pro-preview", "gemini-3-flash-preview", "gemini-3.1-flash-lite-preview",
"gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite",
"gemini-1.5-pro", "gemini-1.5-flash"
);
public String discoverBestModel() {
String apiKey = configService.getGeminiApiKey();
if (apiKey == null || apiKey.trim().isEmpty()) return "None";
try {
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create("https://generativelanguage.googleapis.com/v1beta/models?key=" + apiKey))
.GET().build();
HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
JSONObject json = new JSONObject(response.body());
JSONArray models = json.getJSONArray("models");
List<String> available = new ArrayList<>();
for (int i = 0; i < models.length(); i++) available.add(models.getJSONObject(i).getString("name").replace("models/", ""));
for (String preferred : MODEL_PRIORITY) {
if (available.contains(preferred)) { activeModel = preferred; return activeModel; }
}
}
} catch (Exception e) {}
return activeModel;
}
public String discoverBestRepo(String modelName, String fileName, String metadataContext) {
String apiKey = configService.getGeminiApiKey();
if (apiKey == null || apiKey.trim().isEmpty()) return null;
try {
String shortContext = metadataContext.length() > 25000 ? metadataContext.substring(0, 25000) : metadataContext;
String prompt = "Websuche: Ermittle das offizielle Hugging Face Repository für die Datei '" + modelName + "'.\n\n" +
"KONTEXT:\n" +
"Workflow-Datei: " + fileName + "\n" +
"Workflow-Daten: " + shortContext + "\n\n" +
"ANWEISUNG:\n" +
"1. Identifiziere das exakte Hugging Face Repository (z.B. black-forest-labs/FLUX.1-schnell).\n" +
"2. Antworte NUR mit der Repository ID (Format: creator/repo) oder 'UNKNOWN'.\n" +
"3. Wenn du einen direkten Download-Link findest, gib diesen stattdessen aus.";
JSONObject payload = new JSONObject();
JSONArray contents = new JSONArray();
contents.put(new JSONObject().put("role", "user")
.put("parts", new JSONArray().put(new JSONObject().put("text", prompt))));
payload.put("contents", contents);
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create("https://generativelanguage.googleapis.com/v1beta/models/" + activeModel + ":generateContent?key=" + apiKey))
.header("Content-Type", "application/json")
.POST(HttpRequest.BodyPublishers.ofString(payload.toString()))
.build();
HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
String result = new JSONObject(response.body()).getJSONArray("candidates")
.getJSONObject(0).getJSONObject("content").getJSONArray("parts")
.getJSONObject(0).getString("text").trim();
return result;
}
} catch (Exception e) {}
return null;
}
public String analyzeModel(String modelName) {
String apiKey = configService.getGeminiApiKey();
if (apiKey == null || apiKey.isEmpty()) return null;
try {
JSONObject payload = new JSONObject();
JSONArray contents = new JSONArray();
contents.put(new JSONObject().put("role", "user")
.put("parts", new JSONArray().put(new JSONObject().put("text", "Analyze: " + modelName + ". Return 'Creator | Arch'."))));
payload.put("contents", contents);
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create("https://generativelanguage.googleapis.com/v1beta/models/" + activeModel + ":generateContent?key=" + apiKey))
.header("Content-Type", "application/json")
.POST(HttpRequest.BodyPublishers.ofString(payload.toString()))
.build();
HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
return new JSONObject(response.body()).getJSONArray("candidates").getJSONObject(0).getJSONObject("content").getJSONArray("parts").getJSONObject(0).getString("text").trim();
}
} catch (Exception ignored) {}
return null;
}
}
src/main/java/de/tki/comfymodels/service/impl/LocalAIService.java
+91
View File
@@ -0,0 +1,91 @@
package de.tki.comfymodels.service.impl;
import org.springframework.stereotype.Service;
import java.util.*;
@Service
public class LocalAIService {
private static final Map<String, String[]> KNOWLEDGE_BASE = new LinkedHashMap<>();
private static final Map<String, Double> GLOBAL_IDF = new HashMap<>();
static {
// --- TOP RESEARCH LABS ---
add("black-forest-labs", "flux, flux1, flux2, bfl, schnell, dev, pro, mistral");
add("stabilityai", "sdxl, sd15, sd21, sd3, svd, cascade, turbo, lightning, stable, diffusion");
add("Wan-AI", "wan, wan2, wan2.1, wan2.2, ti2v, t2v");
add("nvidia", "cosmos, world, tokenizer");
add("genmo", "mochi");
add("Comfy-Org", "repack, flux_text_encoders, wan_repackaged, mochi_repackaged, lumina");
add("TencentARC", "hunyuan, hyt");
// --- TOP COMMUNITY CREATORS ---
add("Lykon", "dreamshaper, absolute, reality, anytime, alys");
add("PonyDiffusion", "pony, v6, v7, score_9, rating_explicit");
add("XpucT", "deliberate, real-vision");
add("city96", "gguf, quantized");
add("Kijai", "kj, nodes, flux-fp8");
add("6tZ", "aesthetic, flux_lora");
add("Peli86", "realism, enhance");
add("lllyasviel", "controlnet, adapter, annotator, depth, canny, scribble");
calculateIdf();
}
private static void add(String provider, String keywords) {
KNOWLEDGE_BASE.put(provider, keywords.split(", "));
}
private static void calculateIdf() {
Set<String> allTerms = new HashSet<>();
for (String[] terms : KNOWLEDGE_BASE.values()) allTerms.addAll(Arrays.asList(terms));
int docCount = KNOWLEDGE_BASE.size();
for (String term : allTerms) {
long count = KNOWLEDGE_BASE.values().stream().filter(list -> Arrays.asList(list).contains(term)).count();
GLOBAL_IDF.put(term, Math.log((double) docCount / (1.0 + count)));
}
}
public Prediction predictProvider(String fileName) {
String name = fileName.toLowerCase();
// Spezial-Check für komplexe Dateinamen
if (name.contains("mistral") && name.contains("flux2")) return new Prediction("black-forest-labs", 0.95);
if (name.contains("pony") || name.contains("v6")) return new Prediction("PonyDiffusion", 0.85);
List<String> queryTokens = Arrays.asList(name.split("[_\\- \\.]"));
Map<String, Double> similarityScores = new HashMap<>();
for (Map.Entry<String, String[]> entry : KNOWLEDGE_BASE.entrySet()) {
String provider = entry.getKey();
double score = 0.0;
for (String token : queryTokens) {
if (token.length() < 2) continue;
for (String docTerm : entry.getValue()) {
if (token.equals(docTerm) || (token.contains(docTerm) && docTerm.length() > 3)) {
score += GLOBAL_IDF.getOrDefault(docTerm, 1.0) * (docTerm.length() / 4.0);
}
}
}
if (score > 0) similarityScores.put(provider, score);
}
String bestProvider = "community";
double maxScore = 0.0;
for (Map.Entry<String, Double> entry : similarityScores.entrySet()) {
if (entry.getValue() > maxScore) { maxScore = entry.getValue(); bestProvider = entry.getKey(); }
}
return new Prediction(bestProvider, Math.min(0.99, maxScore / 3.5));
}
public static class Prediction {
public final String provider;
public final double confidence;
public Prediction(String provider, double confidence) { this.provider = provider; this.confidence = confidence; }
public String getLabel() {
if (confidence > 0.75) return "🧠 AI Verified: " + provider;
if (confidence > 0.35) return "🎯 AI Predicted: " + provider;
return "👥 Community";
}
}
}
src/main/java/de/tki/comfymodels/service/impl/ModelListService.java
+110
View File
@@ -0,0 +1,110 @@
package de.tki.comfymodels.service.impl;
import de.tki.comfymodels.domain.ModelInfo;
import org.json.JSONArray;
import org.json.JSONObject;
import org.springframework.stereotype.Service;
import jakarta.annotation.PostConstruct;
import java.io.File;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
@Service
public class ModelListService {
private static final String STORAGE_FILE = "uploaded_models.json";
private List<ModelInfo> models = new ArrayList<>();
@PostConstruct
public void init() {
loadFromStorage();
}
public void importJson(File file) throws Exception {
String content = Files.readString(file.toPath(), StandardCharsets.UTF_8);
JSONObject root = new JSONObject(content);
JSONArray modelsArray = root.getJSONArray("models");
List<ModelInfo> newModels = new ArrayList<>();
for (int i = 0; i < modelsArray.length(); i++) {
JSONObject m = modelsArray.getJSONObject(i);
ModelInfo info = new ModelInfo(
m.optString("type", "Unknown"),
m.optString("name", "Unknown"),
m.optString("url", "MISSING")
);
info.setBase(m.optString("base", ""));
info.setSave_path(m.optString("save_path", ""));
info.setDescription(m.optString("description", ""));
info.setReference(m.optString("reference", ""));
info.setFilename(m.optString("filename", ""));
info.setSize(m.optString("size", "Unknown"));
newModels.add(info);
}
this.models = newModels;
saveToStorage();
}
public Optional<ModelInfo> findByFilename(String filename) {
return models.stream()
.filter(m -> filename.equalsIgnoreCase(m.getFilename()))
.findFirst();
}
private void loadFromStorage() {
try {
File file = new File(STORAGE_FILE);
if (file.exists()) {
String content = Files.readString(file.toPath(), StandardCharsets.UTF_8);
JSONArray array = new JSONArray(content);
for (int i = 0; i < array.length(); i++) {
JSONObject m = array.getJSONObject(i);
ModelInfo info = new ModelInfo(
m.optString("type"),
m.optString("name"),
m.optString("url")
);
info.setBase(m.optString("base"));
info.setSave_path(m.optString("save_path"));
info.setDescription(m.optString("description"));
info.setReference(m.optString("reference"));
info.setFilename(m.optString("filename"));
info.setSize(m.optString("size"));
models.add(info);
}
}
} catch (Exception e) {
e.printStackTrace();
}
}
private void saveToStorage() {
try {
JSONArray array = new JSONArray();
for (ModelInfo m : models) {
JSONObject jo = new JSONObject();
jo.put("type", m.getType());
jo.put("name", m.getName());
jo.put("url", m.getUrl());
jo.put("base", m.getBase());
jo.put("save_path", m.getSave_path());
jo.put("description", m.getDescription());
jo.put("reference", m.getReference());
jo.put("filename", m.getFilename());
jo.put("size", m.getSize());
array.put(jo);
}
Files.writeString(new File(STORAGE_FILE).toPath(), array.toString(4), StandardCharsets.UTF_8);
} catch (Exception e) {
e.printStackTrace();
}
}
public List<ModelInfo> getModels() {
return models;
}
}
src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java
+253
View File
@@ -0,0 +1,253 @@
package de.tki.comfymodels.service.impl;
import de.tki.comfymodels.domain.ModelInfo;
import de.tki.comfymodels.service.IModelSearchService;
import org.json.JSONArray;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import java.net.URI;
import java.net.URLEncoder;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.function.BiConsumer;
@Service
public class ModelSearchService implements IModelSearchService {
@Autowired
private ConfigService configService;
@Autowired
private GeminiAIService geminiService;
@Autowired
private ModelListService modelListService;
private static final Set<String> OFFICIAL_AUTHORS = new HashSet<>(Arrays.asList(
"black-forest-labs", "stabilityai", "runwayml", "comfyanonymous", "Comfy-Org",
"lllyasviel", "Kwai-Kolors", "Kwai-VGI", "InstantX", "ByteDance", "apple",
"google", "facebook", "TencentARC", "DeepFloyd", "microsoft", "nvidia", "fal", "genmo",
"ostris", "XLabs-AI", "THUDM", "city96", "Kijai", "BartoszGawlik"
));
private final HttpClient httpClient = HttpClient.newBuilder()
.followRedirects(HttpClient.Redirect.ALWAYS)
.build();
@Override
public void searchOnline(List<ModelInfo> modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound) {
if (modelsToDownload == null) return;
for (int i = 0; i < modelsToDownload.size(); i++) {
if (!selectedIndices[i]) continue;
ModelInfo info = modelsToDownload.get(i);
final int index = i;
new Thread(() -> {
// Priority 1: User defined Model List
Optional<ModelInfo> manualMatch = modelListService.findByFilename(info.getName());
if (manualMatch.isPresent() && !manualMatch.get().getUrl().equals("MISSING")) {
onStatusUpdate.accept(index, "📂 Found in Model List");
if (validateAndSetUrl(info, index, manualMatch.get().getUrl(), "📂 USER DEFINED", onStatusUpdate, onModelFound)) return;
}
onStatusUpdate.accept(index, "✨ Gemini Scouting...");
String aiHint = geminiService.discoverBestRepo(info.getName(), fileName, workflowContext);
if (aiHint != null && !aiHint.equalsIgnoreCase("UNKNOWN")) {
if (aiHint.startsWith("http")) {
if (validateAndSetUrl(info, index, aiHint, "✨ AI DIRECT", onStatusUpdate, onModelFound)) return;
}
onStatusUpdate.accept(index, "🔍 Validating Repo: " + aiHint);
if (fetchHuggingFaceUrlInSpecificRepo(info, index, aiHint, onStatusUpdate, onModelFound)) return;
}
String modelName = info.getName();
String cleanName = modelName.replaceAll("(_fp8|_fp16|_bf16|_v\\d+|\\d+v|_fix|\\.safetensors|\\.sft|\\.ckpt)", "");
String baseName = cleanName.split("[_\\- ]")[0];
String contextPrefix = fileName.toLowerCase().split("[_\\- ]")[0];
if (contextPrefix.length() < 3) contextPrefix = "";
LinkedHashSet<String> attempts = new LinkedHashSet<>();
if (!contextPrefix.isEmpty()) {
attempts.add(contextPrefix + " " + modelName);
attempts.add(contextPrefix + " " + cleanName);
}
attempts.add(modelName);
attempts.add(cleanName);
attempts.add(baseName);
// Suche in offiziellen Repos (Top 100)
for (String query : attempts) {
if (query.length() < 3) continue;
onStatusUpdate.accept(index, "🔍 Searching Official: " + query);
if (fetchHuggingFaceUrl(info, index, query, true, onStatusUpdate, onModelFound)) return;
}
// Suche in Community Repos (Top 100)
for (String query : attempts) {
if (query.length() < 3) continue;
onStatusUpdate.accept(index, "🔍 Searching Community: " + query);
if (fetchHuggingFaceUrl(info, index, query, false, onStatusUpdate, onModelFound)) return;
}
// Civitai Fallback
for (String query : attempts) {
if (query.length() < 3) continue;
onStatusUpdate.accept(index, "🔍 Searching Civitai: " + query);
if (fetchCivitaiUrl(info, index, query, onStatusUpdate, onModelFound)) return;
}
onStatusUpdate.accept(index, "❌ No trusted match found");
}).start();
}
}
private boolean validateAndSetUrl(ModelInfo info, int rowIndex, String url, String popPrefix,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound) {
long size = getRemoteSize(url);
if (size > 100) {
info.setUrl(url);
info.setPopularity(popPrefix);
info.setSize(formatSize(size));
onModelFound.accept(rowIndex, info);
return true;
}
return false;
}
@Override
public long getRemoteSize(String url) {
try {
if (url.contains("huggingface.co") && url.contains("/resolve/")) {
String apiUrl = url.replace("/resolve/", "/api/models/").replace("/main/", "/file/");
HttpRequest.Builder apiBuilder = HttpRequest.newBuilder().uri(URI.create(apiUrl)).GET().header("User-Agent", "Mozilla/5.0");
String token = configService.getHfToken();
if (!token.isEmpty()) apiBuilder.header("Authorization", "Bearer " + token);
HttpResponse<String> apiRes = httpClient.send(apiBuilder.build(), HttpResponse.BodyHandlers.ofString());
if (apiRes.statusCode() == 200) {
JSONObject fileData = new JSONObject(apiRes.body());
if (fileData.has("size")) return fileData.getLong("size");
}
}
HttpRequest.Builder builder = HttpRequest.newBuilder().uri(URI.create(url)).method("HEAD", HttpRequest.BodyPublishers.noBody()).header("User-Agent", "Mozilla/5.0");
String token = configService.getHfToken();
if (url.contains("huggingface.co") && !token.isEmpty()) builder.header("Authorization", "Bearer " + token);
HttpResponse<Void> res = httpClient.send(builder.build(), HttpResponse.BodyHandlers.discarding());
if (res.statusCode() == 401 || res.statusCode() == 403) return -401;
long bytes = res.headers().firstValueAsLong("Content-Length").orElse(-1L);
if (bytes <= 2000) {
HttpRequest.Builder rBuilder = HttpRequest.newBuilder().uri(URI.create(url)).header("Range", "bytes=0-0").header("User-Agent", "Mozilla/5.0");
if (url.contains("huggingface.co") && !token.isEmpty()) rBuilder.header("Authorization", "Bearer " + token);
HttpResponse<Void> rRes = httpClient.send(rBuilder.build(), HttpResponse.BodyHandlers.discarding());
String cr = rRes.headers().firstValue("Content-Range").orElse("");
bytes = cr.contains("/") ? Long.parseLong(cr.substring(cr.lastIndexOf("/") + 1)) : rRes.headers().firstValueAsLong("Content-Length").orElse(-1L);
}
return bytes;
} catch (Exception e) {
return -1;
}
}
@Override
public String formatSize(long bytes) {
if (bytes == -401) return "🔒 Auth Required";
if (bytes <= 0) return "Unknown";
if (bytes > 1024L * 1024L * 1024L) return String.format("%.2f GB", bytes / 1024.0 / 1024.0 / 1024.0);
return String.format("%.2f MB", bytes / 1024.0 / 1024.0);
}
private boolean fetchHuggingFaceUrlInSpecificRepo(ModelInfo info, int rowIndex, String repoId,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound) {
try {
String treeUrl = "https://huggingface.co/api/models/" + repoId + "/tree/main?recursive=true";
HttpRequest.Builder builder = HttpRequest.newBuilder().uri(URI.create(treeUrl)).GET();
String token = configService.getHfToken();
if (!token.isEmpty()) builder.header("Authorization", "Bearer " + token);
HttpResponse<String> response = httpClient.send(builder.build(), HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
JSONArray files = new JSONArray(response.body());
String targetName = info.getName().toLowerCase();
for (int j = 0; j < files.length(); j++) {
JSONObject file = files.getJSONObject(j);
String path = file.optString("path");
String type = file.optString("type");
if ("file".equals(type)) {
String fileNameOnly = path.contains("/") ? path.substring(path.lastIndexOf("/") + 1) : path;
if (path.equalsIgnoreCase(info.getName()) || fileNameOnly.equalsIgnoreCase(targetName)) {
String downloadUrl = "https://huggingface.co/" + repoId + "/resolve/main/" + path;
return validateAndSetUrl(info, rowIndex, downloadUrl, "✨ AI RECURSIVE", onStatusUpdate, onModelFound);
}
}
}
}
} catch (Exception e) {}
return false;
}
private boolean fetchHuggingFaceUrl(ModelInfo info, int rowIndex, String query, boolean officialOnly,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound) {
try {
String searchUrl = "https://huggingface.co/api/models?search=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=downloads&direction=-1&limit=100";
HttpResponse<String> response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
JSONArray models = new JSONArray(response.body());
for (int i = 0; i < models.length(); i++) {
String modelId = models.getJSONObject(i).getString("id");
String author = modelId.split("/")[0];
if (officialOnly && !OFFICIAL_AUTHORS.contains(author)) continue;
if (fetchHuggingFaceUrlInSpecificRepo(info, rowIndex, modelId, onStatusUpdate, onModelFound)) return true;
}
}
} catch (Exception e) {}
return false;
}
private boolean fetchCivitaiUrl(ModelInfo info, int rowIndex, String query,
BiConsumer<Integer, String> onStatusUpdate,
BiConsumer<Integer, ModelInfo> onModelFound) {
try {
String searchUrl = "https://civitai.com/api/v1/models?query=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=Most+Downloaded&limit=20";
HttpResponse<String> response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString());
if (response.statusCode() == 200) {
JSONObject json = new JSONObject(response.body());
JSONArray items = json.optJSONArray("items");
if (items != null) {
for (int i = 0; i < items.length(); i++) {
JSONArray versions = items.getJSONObject(i).optJSONArray("modelVersions");
if (versions != null) {
for (int v = 0; v < versions.length(); v++) {
JSONArray files = versions.getJSONObject(v).optJSONArray("files");
if (files != null) {
for (int j = 0; j < files.length(); j++) {
JSONObject file = files.getJSONObject(j);
if (file.optString("name").equalsIgnoreCase(info.getName()) || file.optString("name").toLowerCase().contains(query.toLowerCase())) {
return validateAndSetUrl(info, rowIndex, file.optString("downloadUrl"), "🛡️ Civitai", onStatusUpdate, onModelFound);
}
}
}
}
}
}
}
}
} catch (Exception e) {}
return false;
}
}
src/main/java/de/tki/comfymodels/service/impl/WorkflowService.java
+73
View File
@@ -0,0 +1,73 @@
package de.tki.comfymodels.service.impl;
import de.tki.comfymodels.service.IWorkflowService;
import org.json.JSONObject;
import org.springframework.stereotype.Service;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.zip.InflaterInputStream;
@Service
public class WorkflowService implements IWorkflowService {
@Override
public String extractWorkflow(File file) throws IOException {
String name = file.getName().toLowerCase();
if (name.endsWith(".png")) {
return extractWorkflowFromPng(file);
} else {
return Files.readString(file.toPath());
}
}
private String extractWorkflowFromPng(File file) {
String workflow = null, prompt = null;
try (DataInputStream dis = new DataInputStream(new FileInputStream(file))) {
if (dis.readLong() != 0x89504E470D0A1A0AL) return null;
while (true) {
int length = dis.readInt();
byte[] typeBytes = new byte[4];
dis.readFully(typeBytes);
String type = new String(typeBytes, StandardCharsets.US_ASCII);
if ("tEXt".equals(type) || "iTXt".equals(type) || "zTXt".equals(type)) {
byte[] data = new byte[length];
dis.readFully(data);
int nullPos = 0; while (nullPos < data.length && data[nullPos] != 0) nullPos++;
String key = new String(data, 0, nullPos, StandardCharsets.UTF_8);
String value = null;
if ("tEXt".equals(type)) value = new String(data, nullPos + 1, data.length - (nullPos + 1), StandardCharsets.UTF_8);
else if ("zTXt".equals(type)) {
if (nullPos + 2 < data.length) {
try (InflaterInputStream iis = new InflaterInputStream(new ByteArrayInputStream(data, nullPos + 2, data.length - (nullPos + 2)))) { value = new String(iis.readAllBytes(), StandardCharsets.UTF_8); }
}
} else if ("iTXt".equals(type)) {
int currentPos = nullPos + 3; int nullCount = 0;
while (currentPos < data.length && nullCount < 2) { if (data[currentPos] == 0) nullCount++; currentPos++; }
if (currentPos < data.length) {
if (data[nullPos + 1] != 0) {
try (InflaterInputStream iis = new InflaterInputStream(new ByteArrayInputStream(data, currentPos, data.length - currentPos))) { value = new String(iis.readAllBytes(), StandardCharsets.UTF_8); }
} else value = new String(data, currentPos, data.length - currentPos, StandardCharsets.UTF_8);
}
}
if (value != null && !value.isEmpty()) {
if (key.equalsIgnoreCase("workflow")) workflow = value;
else if (key.equalsIgnoreCase("prompt")) prompt = value;
else if (key.toLowerCase().contains("comfy") && workflow == null) workflow = value;
}
dis.readInt();
} else if ("IEND".equals(type)) break; else dis.skipBytes(length + 4);
}
} catch (Exception e) {
// Log error or rethrow as needed
}
if (workflow == null && prompt != null && prompt.trim().startsWith("{")) {
try {
JSONObject jo = new JSONObject(prompt);
if (jo.has("workflow")) return jo.get("workflow").toString();
} catch (Exception ignored) {}
}
return workflow != null ? workflow : prompt;
}
}
src/main/resources/known_good_urls.json
+30
View File
@@ -0,0 +1,30 @@
{
"flux1-schnell.sft": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors",
"flux1-dev.sft": "https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors",
"flux1-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors",
"ae.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors",
"ae.sft": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors",
"flux2-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.2-dev/resolve/main/flux2-dev.safetensors",
"clip_l.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors",
"t5xxl_fp16.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp16.safetensors",
"t5xxl_fp8_e4m3fn.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors",
"clip_g.safetensors": "https://huggingface.co/Comfy-Org/stable-diffusion-v3-medium/resolve/main/text_encoders/clip_g.safetensors",
"v1-5-pruned-emaonly.safetensors": "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors",
"sd_xl_base_1.0.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors",
"sd_xl_refiner_1.0.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/resolve/main/sd_xl_refiner_1.0.safetensors",
"sd3.5_large.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/sd3.5_large.safetensors",
"sd3.5_medium.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-3.5-medium/resolve/main/sd3.5_medium.safetensors",
"sd_xl_turbo_1.0_fp16.safetensors": "https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/sd_xl_turbo_1.0_fp16.safetensors",
"svd.safetensors": "https://huggingface.co/stabilityai/stable-video-diffusion-img2vid/resolve/main/svd.safetensors",
"svd_xt.safetensors": "https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt/resolve/main/svd_xt.safetensors",
"flux1-canny-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev/resolve/main/flux1-canny-dev.safetensors",
"flux1-depth-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev/resolve/main/flux1-depth-dev.safetensors",
"flux1-redux-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev/resolve/main/flux1-redux-dev.safetensors",
"hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors": "https://huggingface.co/Kwai-VGI/HunyuanVideo/resolve/main/hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors",
"mochi_1_preview_fp8.safetensors": "https://huggingface.co/genmo/mochi-1-preview/resolve/main/mochi_1_preview_fp8.safetensors",
"ltx-video-2b-v0.9.safetensors": "https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltx-video-2b-v0.9.safetensors",
"vae-ft-mse-840000-ema-pruned.safetensors": "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors",
"sdxl_vae.safetensors": "https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl_vae.safetensors",
"RealESRGAN_x4plus.pth": "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4plus.pth",
"RealESRGAN_x4plus_anime_6B.pth": "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4plus_anime_6B.pth"
}
src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java
+46
View File
@@ -0,0 +1,46 @@
package de.tki.comfymodels.service;
import de.tki.comfymodels.domain.ModelInfo;
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
import org.junit.jupiter.api.Test;
import java.util.List;
import static org.junit.jupiter.api.Assertions.*;
public class ComfyModelAnalyzerTest {
private final ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
@Test
public void testAnalyzeBasicWorkflow() {
String json = "{\"nodes\": [{\"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"flux1-schnell.sft\"]}]}";
List<ModelInfo> models = analyzer.analyze(json, "test.json");
assertNotNull(models);
assertFalse(models.isEmpty());
ModelInfo flux = models.stream()
.filter(m -> m.getName().equals("flux1-schnell.sft"))
.findFirst()
.orElse(null);
assertNotNull(flux);
assertEquals("checkpoints", flux.getType());
assertEquals("https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors", flux.getUrl());
}
@Test
public void testAnalyzeLora() {
String json = "{\"nodes\": [{\"type\": \"LoraLoader\", \"widgets_values\": [\"my_cool_lora.safetensors\"]}]}";
List<ModelInfo> models = analyzer.analyze(json, "test.json");
ModelInfo lora = models.stream()
.filter(m -> m.getName().equals("my_cool_lora.safetensors"))
.findFirst()
.orElse(null);
assertNotNull(lora);
assertEquals("loras", lora.getType());
}
}
src/test/java/de/tki/comfymodels/WorkflowServiceTest.java
+26
View File
@@ -0,0 +1,26 @@
package de.tki.comfymodels.service;
import de.tki.comfymodels.service.impl.WorkflowService;
import org.junit.jupiter.api.Test;
import java.io.File;
import java.io.IOException;
import static org.junit.jupiter.api.Assertions.*;
public class WorkflowServiceTest {
private final WorkflowService workflowService = new WorkflowService();
@Test
public void testExtractFromPng() throws IOException {
File pngFile = new File("flux_canny_model_example.png");
if (pngFile.exists()) {
String workflow = workflowService.extractWorkflow(pngFile);
assertNotNull(workflow);
assertTrue(workflow.contains("{"));
} else {
System.out.println("Skipping test: flux_canny_model_example.png not found");
}
}
}