From 9d722bf251bdcb70ba48ebb3ab250d124a2cf4dc Mon Sep 17 00:00:00 2001 From: Thomas Date: Sun, 19 Apr 2026 11:30:23 +0200 Subject: [PATCH] Refactor package structure to de.tki.comfymodels, fix compilation errors, and prioritize imported model list in search --- .gitignore | 22 +- pom.xml | 26 +- .../java/com/example/comfymodels/Main.java | 499 ------------------ .../de/tki/comfymodels/BatchVerification.java | 88 +++ src/main/java/de/tki/comfymodels/Main.java | 324 ++++++++++++ src/main/java/de/tki/comfymodels/PngTest.java | 49 ++ .../de/tki/comfymodels/domain/ModelInfo.java | 56 ++ .../comfymodels/service/IDownloadManager.java | 12 + .../comfymodels/service/IModelAnalyzer.java | 8 + .../service/IModelSearchService.java | 23 + .../comfymodels/service/IWorkflowService.java | 11 + .../service/impl/ComfyModelAnalyzer.java | 192 +++++++ .../service/impl/ConfigService.java | 45 ++ .../service/impl/DefaultDownloadManager.java | 121 +++++ .../service/impl/GeminiAIService.java | 120 +++++ .../service/impl/LocalAIService.java | 91 ++++ .../service/impl/ModelListService.java | 110 ++++ .../service/impl/ModelSearchService.java | 253 +++++++++ .../service/impl/WorkflowService.java | 73 +++ src/main/resources/known_good_urls.json | 30 ++ .../comfymodels/ComfyModelAnalyzerTest.java | 46 ++ .../tki/comfymodels/WorkflowServiceTest.java | 26 + 22 files changed, 1715 insertions(+), 510 deletions(-) delete mode 100644 src/main/java/com/example/comfymodels/Main.java create mode 100644 src/main/java/de/tki/comfymodels/BatchVerification.java create mode 100644 src/main/java/de/tki/comfymodels/Main.java create mode 100644 src/main/java/de/tki/comfymodels/PngTest.java create mode 100644 src/main/java/de/tki/comfymodels/domain/ModelInfo.java create mode 100644 src/main/java/de/tki/comfymodels/service/IDownloadManager.java create mode 100644 src/main/java/de/tki/comfymodels/service/IModelAnalyzer.java create mode 100644 src/main/java/de/tki/comfymodels/service/IModelSearchService.java create mode 100644 src/main/java/de/tki/comfymodels/service/IWorkflowService.java create mode 100644 src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java create mode 100644 src/main/java/de/tki/comfymodels/service/impl/ConfigService.java create mode 100644 src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java create mode 100644 src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java create mode 100644 src/main/java/de/tki/comfymodels/service/impl/LocalAIService.java create mode 100644 src/main/java/de/tki/comfymodels/service/impl/ModelListService.java create mode 100644 src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java create mode 100644 src/main/java/de/tki/comfymodels/service/impl/WorkflowService.java create mode 100644 src/main/resources/known_good_urls.json create mode 100644 src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java create mode 100644 src/test/java/de/tki/comfymodels/WorkflowServiceTest.java 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"); + } + } +}