diff --git a/.gitignore b/.gitignore
index 3221cae..0243213 100644
--- a/.gitignore
+++ b/.gitignore
@@ -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
diff --git a/pom.xml b/pom.xml
index 51232f7..a79d352 100644
--- a/pom.xml
+++ b/pom.xml
@@ -4,7 +4,7 @@
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
4.0.0
- com.example
+ de.tki
comfymodeldownloader
1.0-SNAPSHOT
@@ -20,6 +20,28 @@
json
20240303
+
+ org.springframework
+ spring-context
+ 6.1.5
+
+
+ jakarta.annotation
+ jakarta.annotation-api
+ 2.1.1
+
+
+ org.junit.jupiter
+ junit-jupiter-api
+ 5.10.2
+ test
+
+
+ org.junit.jupiter
+ junit-jupiter-engine
+ 5.10.2
+ test
+
@@ -42,7 +64,7 @@
- com.example.comfymodels.Main
+ de.tki.comfymodels.Main
diff --git a/src/main/java/com/example/comfymodels/Main.java b/src/main/java/com/example/comfymodels/Main.java
deleted file mode 100644
index 543f627..0000000
--- a/src/main/java/com/example/comfymodels/Main.java
+++ /dev/null
@@ -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 modelsToDownload = new ArrayList<>();
- private final Set 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 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 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;
- }
- }
-}
diff --git a/src/main/java/de/tki/comfymodels/BatchVerification.java b/src/main/java/de/tki/comfymodels/BatchVerification.java
new file mode 100644
index 0000000..4af0241
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/BatchVerification.java
@@ -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 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;
+ }
+}
diff --git a/src/main/java/de/tki/comfymodels/Main.java b/src/main/java/de/tki/comfymodels/Main.java
new file mode 100644
index 0000000..8387e68
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/Main.java
@@ -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 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 files = (List) 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 {}
+}
diff --git a/src/main/java/de/tki/comfymodels/PngTest.java b/src/main/java/de/tki/comfymodels/PngTest.java
new file mode 100644
index 0000000..0864b40
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/PngTest.java
@@ -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;
+ }
+}
diff --git a/src/main/java/de/tki/comfymodels/domain/ModelInfo.java b/src/main/java/de/tki/comfymodels/domain/ModelInfo.java
new file mode 100644
index 0000000..6d86282
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/domain/ModelInfo.java
@@ -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);
+ }
+}
diff --git a/src/main/java/de/tki/comfymodels/service/IDownloadManager.java b/src/main/java/de/tki/comfymodels/service/IDownloadManager.java
new file mode 100644
index 0000000..a65db5a
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/IDownloadManager.java
@@ -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 models, String baseDir, BiConsumer statusUpdater, Runnable onFinished);
+ void togglePause();
+ void stop();
+ boolean isPaused();
+}
diff --git a/src/main/java/de/tki/comfymodels/service/IModelAnalyzer.java b/src/main/java/de/tki/comfymodels/service/IModelAnalyzer.java
new file mode 100644
index 0000000..b8bb504
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/IModelAnalyzer.java
@@ -0,0 +1,8 @@
+package de.tki.comfymodels.service;
+
+import de.tki.comfymodels.domain.ModelInfo;
+import java.util.List;
+
+public interface IModelAnalyzer {
+ List analyze(String jsonText, String fileName);
+}
diff --git a/src/main/java/de/tki/comfymodels/service/IModelSearchService.java b/src/main/java/de/tki/comfymodels/service/IModelSearchService.java
new file mode 100644
index 0000000..f068005
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/IModelSearchService.java
@@ -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 modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName,
+ BiConsumer onStatusUpdate,
+ BiConsumer onModelFound);
+
+ long getRemoteSize(String url);
+ String formatSize(long bytes);
+}
diff --git a/src/main/java/de/tki/comfymodels/service/IWorkflowService.java b/src/main/java/de/tki/comfymodels/service/IWorkflowService.java
new file mode 100644
index 0000000..008eb9a
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/IWorkflowService.java
@@ -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;
+}
diff --git a/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java b/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java
new file mode 100644
index 0000000..db9c2e5
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java
@@ -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 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 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 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 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 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 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 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 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 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);
+ }
+}
diff --git a/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java b/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java
new file mode 100644
index 0000000..b0f2c7e
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java
@@ -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(); }
+}
diff --git a/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java b/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java
new file mode 100644
index 0000000..56055c2
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java
@@ -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 models, String baseDir, BiConsumer 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 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 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; }
+}
diff --git a/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java b/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java
new file mode 100644
index 0000000..ed9fa4e
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java
@@ -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 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 response = client.send(request, HttpResponse.BodyHandlers.ofString());
+ if (response.statusCode() == 200) {
+ JSONObject json = new JSONObject(response.body());
+ JSONArray models = json.getJSONArray("models");
+ List 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 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 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;
+ }
+}
diff --git a/src/main/java/de/tki/comfymodels/service/impl/LocalAIService.java b/src/main/java/de/tki/comfymodels/service/impl/LocalAIService.java
new file mode 100644
index 0000000..ccb4770
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/impl/LocalAIService.java
@@ -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 KNOWLEDGE_BASE = new LinkedHashMap<>();
+ private static final Map 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 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 queryTokens = Arrays.asList(name.split("[_\\- \\.]"));
+ Map similarityScores = new HashMap<>();
+
+ for (Map.Entry 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 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";
+ }
+ }
+}
diff --git a/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java b/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java
new file mode 100644
index 0000000..6cddcb6
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java
@@ -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 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 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 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 getModels() {
+ return models;
+ }
+}
diff --git a/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java b/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java
new file mode 100644
index 0000000..c028f26
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java
@@ -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 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 modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName,
+ BiConsumer onStatusUpdate,
+ BiConsumer 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 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 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 onStatusUpdate,
+ BiConsumer 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 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 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 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 onStatusUpdate,
+ BiConsumer 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 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 onStatusUpdate,
+ BiConsumer onModelFound) {
+ try {
+ String searchUrl = "https://huggingface.co/api/models?search=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=downloads&direction=-1&limit=100";
+ HttpResponse 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 onStatusUpdate,
+ BiConsumer onModelFound) {
+ try {
+ String searchUrl = "https://civitai.com/api/v1/models?query=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=Most+Downloaded&limit=20";
+ HttpResponse 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;
+ }
+}
diff --git a/src/main/java/de/tki/comfymodels/service/impl/WorkflowService.java b/src/main/java/de/tki/comfymodels/service/impl/WorkflowService.java
new file mode 100644
index 0000000..3fc44ca
--- /dev/null
+++ b/src/main/java/de/tki/comfymodels/service/impl/WorkflowService.java
@@ -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;
+ }
+}
diff --git a/src/main/resources/known_good_urls.json b/src/main/resources/known_good_urls.json
new file mode 100644
index 0000000..387ec75
--- /dev/null
+++ b/src/main/resources/known_good_urls.json
@@ -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"
+}
diff --git a/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java b/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java
new file mode 100644
index 0000000..d269745
--- /dev/null
+++ b/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java
@@ -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 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 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());
+ }
+}
diff --git a/src/test/java/de/tki/comfymodels/WorkflowServiceTest.java b/src/test/java/de/tki/comfymodels/WorkflowServiceTest.java
new file mode 100644
index 0000000..23700db
--- /dev/null
+++ b/src/test/java/de/tki/comfymodels/WorkflowServiceTest.java
@@ -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");
+ }
+ }
+}