feat: add markdown link extraction and HF URL normalization in ComfyModelAnalyzer and ModelSearchService
This commit is contained in:
Thomas Kippster
@@ -18,6 +18,7 @@ Thumbs.db
|
||||
# --- Project Specific ---
|
||||
# Sensitive API keys
|
||||
app_settings.json
|
||||
settings.vault
|
||||
|
||||
# User Model List (internal DB)
|
||||
uploaded_models.json
|
||||
|
||||
@@ -4,10 +4,15 @@ A robust, feature-rich Java Swing application designed to automate the process o
|
||||
|
||||
## 🚀 Key Features
|
||||
|
||||
### 🔍 Advanced JSON Parsing
|
||||
### 🔍 Advanced Workflow Extraction
|
||||
- **JSON & PNG Support**: Directly load `.json` workflows or drop ComfyUI-generated `.png` files to extract embedded workflow and prompt metadata (tEXt, iTXt, zTXt chunks).
|
||||
- **Recursive Traversal**: Scans deeply nested structures, including subgraphs and custom node definitions.
|
||||
- **Filename Detection**: Uses smart regex to find model files (`.safetensors`, `.ckpt`, `.pt`, etc.) inside `widgets_values`, `inputs`, and other string fields.
|
||||
- **Metadata Support**: Automatically detects and prioritizes embedded model metadata (standard `models` array used by ComfyUI managers).
|
||||
- **Smart Detection**: Uses regex and node-type analysis to find model files (`.safetensors`, `.ckpt`, etc.) and infer their target directory (e.g., `checkpoints`, `loras`, `vae`).
|
||||
|
||||
### 🤖 AI-Powered Discovery
|
||||
- **Local AI Heuristics**: Built-in scoring engine to predict model providers (e.g., StabilityAI, Black Forest Labs) offline based on filenames.
|
||||
- **Gemini Deep Search**: Integrated "Deep Search" using the Google Gemini API to find official Hugging Face repositories and direct download links for obscure models.
|
||||
- **Context Awareness**: Disambiguates generic filenames (like `ae.safetensors`) by analyzing the global workflow context (e.g., detecting FLUX or SD3 usage).
|
||||
|
||||
### 📥 Powerful Download Manager
|
||||
- **Resume Capability**: Implements HTTP Range requests to continue interrupted downloads exactly where they left off.
|
||||
@@ -22,38 +27,30 @@ A robust, feature-rich Java Swing application designed to automate the process o
|
||||
## 🛠️ Requirements
|
||||
- **Java 11** or higher
|
||||
- **Maven** (for building from source)
|
||||
- **API Key (Optional)**: A Google Gemini API key for "Deep Search" functionality.
|
||||
|
||||
## 🏗️ Installation & Build
|
||||
## 🏗️ Build & Run
|
||||
|
||||
1. Clone the repository:
|
||||
```bash
|
||||
git clone https://github.com/thomaskippster/comfymodeldownloader.git
|
||||
cd comfymodeldownloader
|
||||
```
|
||||
|
||||
2. Build the "Fat JAR" using Maven:
|
||||
1. **Build the Fat JAR**:
|
||||
```bash
|
||||
mvn clean package
|
||||
```
|
||||
|
||||
3. Locate the runnable JAR in the `target/` directory:
|
||||
`comfymodeldownloader-1.0-SNAPSHOT.jar`
|
||||
2. **Launch**:
|
||||
```bash
|
||||
java -jar target/comfymodeldownloader-1.0-SNAPSHOT.jar
|
||||
```
|
||||
|
||||
## 📖 How to Use
|
||||
|
||||
1. **Launch**: Start the application via `java -jar target/comfymodeldownloader-1.0-SNAPSHOT.jar`.
|
||||
2. **Setup Path**: Enter or browse to your local ComfyUI models directory (e.g., `C:\ComfyUI\models`).
|
||||
3. **Provide Workflow**:
|
||||
- **Paste**: Directly paste the content of your ComfyUI JSON file into the text area.
|
||||
- **Load**: Use the "Load File..." button to select a `.json` file from your disk.
|
||||
4. **Analyze**: Click **Analyze Models**. The table will populate with detected models and their types.
|
||||
5. **Download**:
|
||||
- Select models using the checkboxes.
|
||||
- Click **Start Queue**.
|
||||
- If a file exists, choose to **Resume** or **Overwrite**.
|
||||
1. **Setup Path**: Enter or browse to your local ComfyUI models directory.
|
||||
2. **Provide Workflow**: Paste JSON content, load a `.json` file, or load a `.png` with embedded metadata.
|
||||
3. **Analyze**: Click **Analyze Models**. The application will identify missing models and categorize them.
|
||||
4. **Deep Search**: For unknown models, use the "Deep Search" feature to let the AI find the source.
|
||||
5. **Download**: Select models and click **Start Queue**.
|
||||
|
||||
## 📜 License
|
||||
This project is licensed under the MIT License - see the LICENSE file for details (or just use it freely!).
|
||||
This project is licensed under the MIT License.
|
||||
|
||||
---
|
||||
*Developed to make ComfyUI workflow sharing seamless.*
|
||||
|
||||
@@ -5,12 +5,10 @@
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
|
||||
<groupId>de.tki</groupId>
|
||||
<artifactId>comfymodeldownloader</artifactId>
|
||||
<artifactId>comfyuimodel-downloader</artifactId>
|
||||
<version>1.0-SNAPSHOT</version>
|
||||
|
||||
<properties>
|
||||
<maven.compiler.source>11</maven.compiler.source>
|
||||
<maven.compiler.target>11</maven.compiler.target>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
</properties>
|
||||
|
||||
@@ -42,6 +40,12 @@
|
||||
<version>5.10.2</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.springframework</groupId>
|
||||
<artifactId>spring-test</artifactId>
|
||||
<version>6.1.5</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
@@ -50,6 +54,9 @@
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-compiler-plugin</artifactId>
|
||||
<version>3.11.0</version>
|
||||
<configuration>
|
||||
<release>11</release>
|
||||
</configuration>
|
||||
</plugin>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
@echo off
|
||||
echo Building the project with Maven...
|
||||
call mvn clean package
|
||||
if %ERRORLEVEL% NEQ 0 (
|
||||
echo Build failed! Please check the errors above.
|
||||
pause
|
||||
exit /b %ERRORLEVEL%
|
||||
)
|
||||
|
||||
echo.
|
||||
echo Starting ComfyUI Model Downloader...
|
||||
java -jar target\comfyuimodel-downloader-1.0-SNAPSHOT.jar
|
||||
pause
|
||||
@@ -16,20 +16,17 @@ import org.springframework.stereotype.Component;
|
||||
|
||||
import javax.swing.*;
|
||||
import javax.swing.table.DefaultTableModel;
|
||||
import javax.swing.event.TableModelListener;
|
||||
import javax.swing.event.TableModelEvent;
|
||||
import java.awt.*;
|
||||
import java.awt.datatransfer.DataFlavor;
|
||||
import java.awt.dnd.*;
|
||||
import java.awt.event.MouseAdapter;
|
||||
import java.awt.event.MouseEvent;
|
||||
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 {
|
||||
@@ -37,7 +34,7 @@ public class Main extends JFrame {
|
||||
private final IDownloadManager downloadManager;
|
||||
private final IWorkflowService workflowService;
|
||||
private final IModelSearchService searchService;
|
||||
|
||||
|
||||
@Autowired
|
||||
private ConfigService configService;
|
||||
|
||||
@@ -53,13 +50,12 @@ public class Main extends JFrame {
|
||||
private JLabel activeAiModelLabel;
|
||||
private JTextArea jsonInputArea;
|
||||
private DefaultTableModel tableModel;
|
||||
private JProgressBar progressBar;
|
||||
private JLabel statusLabel;
|
||||
private JButton downloadButton, pauseButton, stopButton;
|
||||
private List<ModelInfo> modelsToDownload;
|
||||
private String currentFileName = "input.json";
|
||||
|
||||
public Main(IModelAnalyzer analyzer, IDownloadManager downloadManager,
|
||||
public Main(IModelAnalyzer analyzer, IDownloadManager downloadManager,
|
||||
IWorkflowService workflowService, IModelSearchService searchService) {
|
||||
this.analyzer = analyzer;
|
||||
this.downloadManager = downloadManager;
|
||||
@@ -67,11 +63,77 @@ public class Main extends JFrame {
|
||||
this.searchService = searchService;
|
||||
}
|
||||
|
||||
@Autowired
|
||||
public void start() {
|
||||
initUI();
|
||||
updateAiModelDisplay();
|
||||
SwingUtilities.invokeLater(() -> setVisible(true));
|
||||
public void launch() {
|
||||
try {
|
||||
UIManager.setLookAndFeel(UIManager.getSystemLookAndFeelClassName());
|
||||
} catch (Exception e) {}
|
||||
|
||||
if (!promptForPassword()) {
|
||||
System.exit(0);
|
||||
}
|
||||
|
||||
SwingUtilities.invokeLater(() -> {
|
||||
initUI();
|
||||
loadSettingsIntoUI();
|
||||
updateAiModelDisplay();
|
||||
|
||||
// Add WindowListener to stop downloads on close
|
||||
addWindowListener(new java.awt.event.WindowAdapter() {
|
||||
@Override
|
||||
public void windowClosing(java.awt.event.WindowEvent e) {
|
||||
downloadManager.stop();
|
||||
System.exit(0);
|
||||
}
|
||||
});
|
||||
|
||||
setVisible(true);
|
||||
});
|
||||
}
|
||||
|
||||
private boolean promptForPassword() {
|
||||
while (true) {
|
||||
JPanel panel = new JPanel(new BorderLayout(5, 5));
|
||||
panel.add(new JLabel("Enter Vault Password (to unlock API Keys):"), BorderLayout.NORTH);
|
||||
JPasswordField pf = new JPasswordField();
|
||||
panel.add(pf, BorderLayout.CENTER);
|
||||
|
||||
// Focus request
|
||||
SwingUtilities.invokeLater(() -> pf.requestFocusInWindow());
|
||||
|
||||
int ok = JOptionPane.showConfirmDialog(null, panel, "Vault Unlock", JOptionPane.OK_CANCEL_OPTION, JOptionPane.QUESTION_MESSAGE);
|
||||
if (ok != JOptionPane.OK_OPTION) return false;
|
||||
|
||||
String pass = new String(pf.getPassword());
|
||||
if (pass.trim().isEmpty()) {
|
||||
JOptionPane.showMessageDialog(null, "Password cannot be empty.", "Validation", JOptionPane.WARNING_MESSAGE);
|
||||
continue;
|
||||
}
|
||||
|
||||
try {
|
||||
configService.unlock(pass);
|
||||
return true;
|
||||
} catch (Exception e) {
|
||||
JOptionPane.showMessageDialog(null, "Unlock Failed: " + e.getMessage(), "Error", JOptionPane.ERROR_MESSAGE);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private void loadSettingsIntoUI() {
|
||||
if (modelsPathField != null) {
|
||||
String path = configService.getModelsPath();
|
||||
modelsPathField.setText(path);
|
||||
System.out.println("GUI: Models Path loaded: " + path);
|
||||
}
|
||||
if (geminiKeyField != null) {
|
||||
String key = configService.getGeminiApiKey();
|
||||
geminiKeyField.setText(key);
|
||||
System.out.println("GUI: Gemini Key loaded (length: " + (key != null ? key.length() : 0) + ")");
|
||||
}
|
||||
if (hfTokenField != null) {
|
||||
String token = configService.getHfToken();
|
||||
hfTokenField.setText(token);
|
||||
System.out.println("GUI: HF Token loaded (length: " + (token != null ? token.length() : 0) + ")");
|
||||
}
|
||||
}
|
||||
|
||||
private void updateAiModelDisplay() {
|
||||
@@ -82,30 +144,48 @@ public class Main extends JFrame {
|
||||
}
|
||||
|
||||
private void initUI() {
|
||||
setTitle("ComfyUI Downloader (ULTRA-VALIDATION)");
|
||||
setTitle("ComfyUIModel-Downloader");
|
||||
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");
|
||||
modelsPathField = new JTextField();
|
||||
JButton browsePathBtn = new JButton("Set Folder...");
|
||||
browsePathBtn.addActionListener(e -> {
|
||||
JFileChooser chooser = new JFileChooser();
|
||||
chooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
|
||||
String current = modelsPathField.getText();
|
||||
if (current != null && !current.isEmpty()) {
|
||||
File dir = new File(current);
|
||||
if (dir.exists()) chooser.setCurrentDirectory(dir);
|
||||
}
|
||||
if (chooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
|
||||
modelsPathField.setText(chooser.getSelectedFile().getAbsolutePath());
|
||||
String selectedPath = chooser.getSelectedFile().getAbsolutePath();
|
||||
modelsPathField.setText(selectedPath);
|
||||
configService.setModelsPath(selectedPath);
|
||||
}
|
||||
});
|
||||
|
||||
JButton savePathBtn = new JButton("Save Path");
|
||||
savePathBtn.addActionListener(e -> {
|
||||
configService.setModelsPath(modelsPathField.getText());
|
||||
statusLabel.setText("Models path saved to vault.");
|
||||
});
|
||||
|
||||
JPanel pathButtons = new JPanel(new FlowLayout(FlowLayout.RIGHT, 0, 0));
|
||||
pathButtons.add(savePathBtn);
|
||||
pathButtons.add(browsePathBtn);
|
||||
|
||||
pathRow.add(modelsPathField, BorderLayout.CENTER);
|
||||
pathRow.add(browsePathBtn, BorderLayout.EAST);
|
||||
pathRow.add(pathButtons, BorderLayout.EAST);
|
||||
|
||||
JPanel apiRow = new JPanel(new BorderLayout());
|
||||
apiRow.setBorder(BorderFactory.createTitledBorder("Gemini AI API Key"));
|
||||
geminiKeyField = new JPasswordField(configService.getGeminiApiKey());
|
||||
geminiKeyField = new JPasswordField();
|
||||
JButton saveGeminiBtn = new JButton("Save & Discover");
|
||||
saveGeminiBtn.addActionListener(e -> {
|
||||
configService.setGeminiApiKey(new String(geminiKeyField.getPassword()));
|
||||
@@ -116,21 +196,20 @@ public class Main extends JFrame {
|
||||
|
||||
JPanel hfRow = new JPanel(new BorderLayout());
|
||||
hfRow.setBorder(BorderFactory.createTitledBorder("Hugging Face Access Token"));
|
||||
hfTokenField = new JPasswordField(configService.getHfToken());
|
||||
hfTokenField = new JPasswordField();
|
||||
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
|
||||
analyzeJsonContent();
|
||||
});
|
||||
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);
|
||||
JButton helpBtn = new JButton("ℹ Help & Security Information");
|
||||
helpBtn.addActionListener(e -> showHelpDialog());
|
||||
infoRow.add(helpBtn);
|
||||
|
||||
settingsPanel.add(pathRow);
|
||||
settingsPanel.add(apiRow);
|
||||
@@ -143,7 +222,7 @@ public class Main extends JFrame {
|
||||
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 -> {
|
||||
@@ -159,34 +238,52 @@ public class Main extends JFrame {
|
||||
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
|
||||
JOptionPane.showMessageDialog(this, "Model list imported successfully!");
|
||||
analyzeJsonContent();
|
||||
} catch (Exception ex) {
|
||||
JOptionPane.showMessageDialog(this, "Error importing model list: " + ex.getMessage());
|
||||
JOptionPane.showMessageDialog(this, "Error: " + 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"};
|
||||
String[] columnNames = {"Select", "Type", "Name", "Size", "AI Source", "Target Path", "URL", "Status"};
|
||||
tableModel = new DefaultTableModel(columnNames, 0) {
|
||||
@Override public Class<?> getColumnClass(int c) { return c == 0 ? Boolean.class : String.class; }
|
||||
@Override public Class<?> getColumnClass(int c) { return c == 0 ? Boolean.class : String.class; }
|
||||
@Override public boolean isCellEditable(int r, int c) { return c == 0; }
|
||||
};
|
||||
tableModel.addTableModelListener(e -> {
|
||||
if (e.getType() == TableModelEvent.UPDATE && e.getColumn() == 0) {
|
||||
updateDownloadManagerSelection();
|
||||
}
|
||||
});
|
||||
JTable modelTable = new JTable(tableModel);
|
||||
modelTable.putClientProperty("terminateEditOnFocusLost", Boolean.TRUE);
|
||||
|
||||
// Sofortiger Abbruch des Editors bei Klick, um Wert zu übernehmen
|
||||
modelTable.addMouseListener(new MouseAdapter() {
|
||||
@Override
|
||||
public void mouseReleased(MouseEvent e) {
|
||||
if (modelTable.isEditing()) {
|
||||
modelTable.getCellEditor().stopCellEditing();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
modelTable.getColumnModel().getColumn(0).setMaxWidth(60);
|
||||
modelTable.getColumnModel().getColumn(4).setMinWidth(250);
|
||||
modelTable.getColumnModel().getColumn(4).setMinWidth(150);
|
||||
modelTable.getColumnModel().getColumn(5).setMinWidth(150);
|
||||
JScrollPane tableScroll = new JScrollPane(modelTable);
|
||||
JPanel tablePanel = new JPanel(new BorderLayout());
|
||||
tablePanel.setBorder(BorderFactory.createTitledBorder("Identified Models"));
|
||||
@@ -197,16 +294,25 @@ public class Main extends JFrame {
|
||||
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 -> {
|
||||
int rowCount = tableModel.getRowCount();
|
||||
boolean[] selected = new boolean[rowCount];
|
||||
for (int i = 0; i < rowCount; i++) {
|
||||
selected[i] = (Boolean) tableModel.getValueAt(i, 0);
|
||||
String urlFromTable = (String) tableModel.getValueAt(i, 6);
|
||||
if (urlFromTable != null && !urlFromTable.equals("MISSING")) {
|
||||
modelsToDownload.get(i).setUrl(urlFromTable);
|
||||
}
|
||||
}
|
||||
|
||||
downloadButton.setEnabled(false);
|
||||
downloadManager.startQueue(modelsToDownload, modelsPathField.getText(),
|
||||
(idx, status) -> SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, idx, 6)),
|
||||
downloadManager.startQueue(modelsToDownload, selected, modelsPathField.getText(),
|
||||
(idx, status) -> SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, idx, 7)),
|
||||
() -> SwingUtilities.invokeLater(() -> {
|
||||
downloadButton.setEnabled(true);
|
||||
statusLabel.setText("Queue finished.");
|
||||
@@ -216,20 +322,26 @@ public class Main extends JFrame {
|
||||
|
||||
pauseButton = new JButton("Pause");
|
||||
pauseButton.addActionListener(e -> {
|
||||
if (downloadManager.isPaused()) {
|
||||
updateDownloadManagerSelection();
|
||||
}
|
||||
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));
|
||||
activeAiModelLabel = new JLabel("Active AI: Detecting...");
|
||||
activeAiModelLabel.setFont(new Font("SansSerif", Font.ITALIC, 11));
|
||||
activeAiModelLabel.setForeground(Color.GRAY);
|
||||
progressPanel.add(activeAiModelLabel);
|
||||
progressPanel.add(statusLabel);
|
||||
progressPanel.add(progressBar);
|
||||
bottomPanel.add(progressPanel, BorderLayout.CENTER);
|
||||
bottomPanel.add(actionButtons, BorderLayout.EAST);
|
||||
|
||||
@@ -238,6 +350,15 @@ public class Main extends JFrame {
|
||||
add(bottomPanel, BorderLayout.SOUTH);
|
||||
}
|
||||
|
||||
private void updateDownloadManagerSelection() {
|
||||
int rowCount = tableModel.getRowCount();
|
||||
boolean[] selected = new boolean[rowCount];
|
||||
for (int i = 0; i < rowCount; i++) {
|
||||
selected[i] = (Boolean) tableModel.getValueAt(i, 0);
|
||||
}
|
||||
downloadManager.updateSelection(selected);
|
||||
}
|
||||
|
||||
private void setupDragAndDrop(JTextArea area) {
|
||||
new DropTarget(area, new DropTargetListener() {
|
||||
public void dragEnter(DropTargetDragEvent dtde) {}
|
||||
@@ -247,7 +368,7 @@ public class Main extends JFrame {
|
||||
public void drop(DropTargetDropEvent dtde) {
|
||||
try {
|
||||
dtde.acceptDrop(DnDConstants.ACTION_COPY);
|
||||
List<File> files = (List<File>) dtde.getTransferable().getTransferData(DataFlavor.javaFileListFlavor);
|
||||
java.util.List<File> files = (java.util.List<File>) dtde.getTransferable().getTransferData(DataFlavor.javaFileListFlavor);
|
||||
if (!files.isEmpty()) loadFile(files.get(0));
|
||||
} catch (Exception e) {}
|
||||
}
|
||||
@@ -270,7 +391,8 @@ public class Main extends JFrame {
|
||||
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"});
|
||||
String targetSubDir = "models" + File.separator + (info.getType() != null ? info.getType() : "checkpoints");
|
||||
tableModel.addRow(new Object[]{true, info.getType(), info.getName(), info.getSize(), info.getPopularity(), targetSubDir, info.getUrl(), info.getUrl().equals("MISSING") ? "Idle" : "✅ Known Good"});
|
||||
if (!info.getUrl().equals("MISSING")) fetchSizeInBackground(info, i);
|
||||
}
|
||||
downloadButton.setEnabled(!modelsToDownload.isEmpty());
|
||||
@@ -278,32 +400,65 @@ public class Main extends JFrame {
|
||||
|
||||
private void searchMissingOnline() {
|
||||
if (modelsToDownload == null) return;
|
||||
|
||||
int rowCount = tableModel.getRowCount();
|
||||
boolean[] selected = new boolean[rowCount];
|
||||
boolean anySelected = false;
|
||||
boolean[] selectedForSearch = new boolean[rowCount];
|
||||
boolean anyMissingSelected = false;
|
||||
for (int i = 0; i < rowCount; i++) {
|
||||
selected[i] = (Boolean) tableModel.getValueAt(i, 0);
|
||||
if (selected[i]) anySelected = true;
|
||||
boolean isChecked = (Boolean) tableModel.getValueAt(i, 0);
|
||||
String currentUrl = (String) tableModel.getValueAt(i, 6);
|
||||
if (isChecked && (currentUrl == null || currentUrl.equals("MISSING") || currentUrl.isEmpty())) {
|
||||
selectedForSearch[i] = true;
|
||||
anyMissingSelected = true;
|
||||
} else {
|
||||
selectedForSearch[i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
if (!anySelected) {
|
||||
JOptionPane.showMessageDialog(this, "Please select at least one model in the table to search.");
|
||||
if (!anyMissingSelected) {
|
||||
JOptionPane.showMessageDialog(this, "No missing models selected for Deep Search.");
|
||||
return;
|
||||
}
|
||||
|
||||
statusLabel.setText("Deep Search in progress...");
|
||||
searchService.searchOnline(modelsToDownload, selected, jsonInputArea.getText(), currentFileName,
|
||||
(idx, status) -> SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, idx, 6)),
|
||||
searchService.searchOnline(modelsToDownload, selectedForSearch, jsonInputArea.getText(), currentFileName,
|
||||
(idx, status) -> SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, idx, 7)),
|
||||
(idx, info) -> SwingUtilities.invokeLater(() -> {
|
||||
String targetSubDir = "models" + File.separator + (info.getType() != null ? info.getType() : "checkpoints");
|
||||
tableModel.setValueAt(info.getSize(), idx, 3);
|
||||
tableModel.setValueAt(info.getPopularity(), idx, 4);
|
||||
tableModel.setValueAt(info.getUrl(), idx, 5);
|
||||
tableModel.setValueAt("✅ Found", idx, 6);
|
||||
})
|
||||
tableModel.setValueAt(targetSubDir, idx, 5);
|
||||
tableModel.setValueAt(info.getUrl(), idx, 6);
|
||||
tableModel.setValueAt("✅ Found", idx, 7);
|
||||
}),
|
||||
() -> SwingUtilities.invokeLater(() -> statusLabel.setText("Deep Search finished."))
|
||||
);
|
||||
}
|
||||
|
||||
private void showHelpDialog() {
|
||||
String helpText = "<html><body style='width: 500px; padding: 10px;'>" +
|
||||
"<h1>ComfyUIModel-Downloader Help</h1>" +
|
||||
"<h3>Core Features</h3>" +
|
||||
"<ul>" +
|
||||
" <li><b>Workflow Analysis:</b> Drag and drop ComfyUI JSON or PNG files to extract required models.</li>" +
|
||||
" <li><b>Deep Search:</b> Uses AI and multi-platform searching (Hugging Face, Civitai) to find missing model URLs.</li>" +
|
||||
" <li><b>Model List Integration:</b> Match models against your own known-good local lists.</li>" +
|
||||
" <li><b>Queue Management:</b> Download multiple models simultaneously with pause/resume support.</li>" +
|
||||
"</ul>" +
|
||||
"<h3>API Keys & Security</h3>" +
|
||||
"<ul>" +
|
||||
" <li><b>Gemini API Key:</b> Required for the 'AI Scouting' feature. Gemini analyzes your workflow context " +
|
||||
" to predict the best repositories for obscure models.</li>" +
|
||||
" <li><b>Hugging Face Token:</b> Recommended to avoid rate limits and to access gated repositories " +
|
||||
" (like some FLUX or StabilityAI models).</li>" +
|
||||
" <li><b>Local Storage (Vault):</b> All keys and settings are stored <b>locally</b> on your machine " +
|
||||
" in an encrypted file (<code>settings.vault</code>). They are never sent to any server except " +
|
||||
" the official Google/Hugging Face APIs during requests.</li>" +
|
||||
" <li><b>Password Protection:</b> Your vault is protected by the password you choose at startup. " +
|
||||
" Without this password, the API keys cannot be decrypted.</li>" +
|
||||
"</ul>" +
|
||||
"</body></html>";
|
||||
|
||||
JOptionPane.showMessageDialog(this, helpText, "Application Help & Security", JOptionPane.INFORMATION_MESSAGE);
|
||||
}
|
||||
|
||||
private void fetchSizeInBackground(ModelInfo info, int rowIndex) {
|
||||
new Thread(() -> {
|
||||
long size = searchService.getRemoteSize(info.getUrl());
|
||||
@@ -314,11 +469,12 @@ public class Main extends JFrame {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(AppConfig.class);
|
||||
context.getBean(Main.class);
|
||||
AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(AppConfig.class);
|
||||
Main main = context.getBean(Main.class);
|
||||
main.launch();
|
||||
}
|
||||
|
||||
@Configuration
|
||||
@ComponentScan("de.tki.comfymodels")
|
||||
public static class AppConfig {}
|
||||
}
|
||||
}
|
||||
@@ -5,7 +5,8 @@ import java.util.List;
|
||||
import java.util.function.BiConsumer;
|
||||
|
||||
public interface IDownloadManager {
|
||||
void startQueue(List<ModelInfo> models, String baseDir, BiConsumer<Integer, String> statusUpdater, Runnable onFinished);
|
||||
void startQueue(List<ModelInfo> models, boolean[] selectedIndices, String baseDir, BiConsumer<Integer, String> statusUpdater, Runnable onFinished);
|
||||
void updateSelection(boolean[] selectedIndices);
|
||||
void togglePause();
|
||||
void stop();
|
||||
boolean isPaused();
|
||||
|
||||
@@ -16,7 +16,8 @@ public interface IModelSearchService {
|
||||
*/
|
||||
void searchOnline(List<ModelInfo> modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName,
|
||||
BiConsumer<Integer, String> onStatusUpdate,
|
||||
BiConsumer<Integer, ModelInfo> onModelFound);
|
||||
BiConsumer<Integer, ModelInfo> onModelFound,
|
||||
Runnable onFinished);
|
||||
|
||||
long getRemoteSize(String url);
|
||||
String formatSize(long bytes);
|
||||
|
||||
@@ -28,10 +28,11 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
|
||||
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 FILE = Pattern.compile("([a-zA-Z0-9_\\-\\.\\/]+\\.(?:safetensors|sft|ckpt|pth|pt|bin|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);
|
||||
static final Pattern COMP_CLIP = Pattern.compile("(clip|t5|text_encoder|llama|gemma|mistral|t5xxl|clip_l|clip_g|qwen)", Pattern.CASE_INSENSITIVE);
|
||||
static final Pattern COMP_UNET = Pattern.compile("(unet|diffusion|model|transformer|base|transformer|upscale|esrgan|z_image|lumina|flux|dit|mochi|ltx|cosmos|hunyuan|hy|wan|hidream|aura|svd)", Pattern.CASE_INSENSITIVE);
|
||||
static final Pattern MD_LINK = Pattern.compile("\\[([^\\]]+)\\]\\((https?://[^\\)]+)\\)", Pattern.CASE_INSENSITIVE);
|
||||
}
|
||||
|
||||
private final Map<String, String> knownGoodUrls = new HashMap<>();
|
||||
@@ -62,14 +63,23 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
|
||||
try {
|
||||
JSONObject jo = parseJson(jsonText);
|
||||
if (jo == null) return results;
|
||||
|
||||
// NEW: Globale Kontext-Erkennung vor der Analyse
|
||||
String globalContext = extractGlobalContext(jo, fileName);
|
||||
|
||||
findModelsMetadata(jo, results);
|
||||
scanForModelFiles(jo, null, results);
|
||||
applyContext(results, fileName);
|
||||
|
||||
// Wende globalen Kontext auf alle Ergebnisse an
|
||||
applyGlobalContext(results, globalContext);
|
||||
|
||||
for (ModelInfo info : results) {
|
||||
String low = info.getName().toLowerCase();
|
||||
Optional<ModelInfo> uploadedMatch = (modelListService != null) ? modelListService.findByFilename(info.getName()) : Optional.empty();
|
||||
if (uploadedMatch.isPresent()) {
|
||||
info.setUrl(uploadedMatch.get().getUrl());
|
||||
ModelInfo match = uploadedMatch.get();
|
||||
info.setUrl(match.getUrl());
|
||||
info.setSave_path(match.getSave_path());
|
||||
info.setPopularity("📂 USER DEFINED LIST");
|
||||
} else if (knownGoodUrls.containsKey(low)) {
|
||||
info.setUrl(knownGoodUrls.get(low));
|
||||
@@ -77,12 +87,56 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
|
||||
} else {
|
||||
LocalAIService.Prediction prediction = aiService.predictProvider(info.getName());
|
||||
info.setPopularity(prediction.getLabel());
|
||||
|
||||
// Falls die Popularität noch generisch ist, füge den Kontext hinzu
|
||||
if (globalContext != null && !globalContext.isEmpty() && info.getPopularity().contains("Community")) {
|
||||
info.setPopularity("✨ Context: " + globalContext + " (" + info.getPopularity() + ")");
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (Exception e) {}
|
||||
return results;
|
||||
}
|
||||
|
||||
private String extractGlobalContext(JSONObject jo, String fileName) {
|
||||
Map<String, Integer> hints = new HashMap<>();
|
||||
|
||||
// 1. Aus Dateiname (höchste Prio)
|
||||
if (fileName != null && !fileName.toLowerCase().startsWith("workflow")) {
|
||||
String base = fileName.split("[_\\- ]")[0].toLowerCase();
|
||||
if (base.length() >= 3) return base;
|
||||
}
|
||||
|
||||
// 2. Scan den ganzen JSON nach signifikanten Schlüsselwörtern
|
||||
String fullText = jo.toString().toLowerCase();
|
||||
String[] keywords = {"hidream", "flux", "auraflow", "sd3", "hunyuan", "mochi", "ltxv", "cosmos", "wan2.1"};
|
||||
for (String kw : keywords) {
|
||||
if (fullText.contains(kw)) hints.put(kw, hints.getOrDefault(kw, 0) + 5);
|
||||
}
|
||||
|
||||
// 3. Scan Node-Typen
|
||||
String nodesText = fullText;
|
||||
if (nodesText.contains("hidream")) hints.put("hidream", hints.getOrDefault("hidream", 0) + 10);
|
||||
|
||||
return hints.entrySet().stream()
|
||||
.max(Map.Entry.comparingByValue())
|
||||
.map(Map.Entry::getKey)
|
||||
.orElse(null);
|
||||
}
|
||||
|
||||
private void applyGlobalContext(List<ModelInfo> res, String ctx) {
|
||||
if (ctx == null) return;
|
||||
for (ModelInfo i : res) {
|
||||
// Wenn der Dateiname den Kontext enthält oder der Dateiname sehr generisch ist (wie ae.safetensors)
|
||||
String n = i.getName().toLowerCase();
|
||||
if (n.contains(ctx) || n.equals("ae.safetensors") || n.contains("clip_l") || n.contains("t5xxl")) {
|
||||
if (i.getPopularity() == null || i.getPopularity().isEmpty()) {
|
||||
i.setPopularity("✨ Context: " + ctx);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private JSONObject parseJson(String text) {
|
||||
try {
|
||||
String trimmed = text.trim();
|
||||
@@ -109,7 +163,13 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
|
||||
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")));
|
||||
if (n != null && !n.isEmpty()) {
|
||||
// Strikte Übernahme des Pfades aus den Metadaten
|
||||
String rawPath = m.optString("directory", m.optString("type", "checkpoints"));
|
||||
ModelInfo info = new ModelInfo(rawPath, n, m.optString("url", "MISSING"));
|
||||
info.setSave_path(rawPath);
|
||||
addModelInfo(res, info);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -142,16 +202,67 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
|
||||
}
|
||||
|
||||
private void processVal(String v, String ctx, String hint, List<ModelInfo> res) {
|
||||
if (v == null || v.trim().isEmpty()) return;
|
||||
String vLow = v.toLowerCase().trim();
|
||||
|
||||
// Ignoriere Standardwerte, die keine Dateinamen sind
|
||||
if (vLow.equals("none") || vLow.equals("default") || vLow.equals("undefined") || vLow.length() < 4) return;
|
||||
|
||||
// 1. Suche nach Markdown Links: [Name](URL)
|
||||
Matcher mdMatcher = ExpertPatterns.MD_LINK.matcher(v);
|
||||
while (mdMatcher.find()) {
|
||||
String name = mdMatcher.group(1);
|
||||
String url = mdMatcher.group(2);
|
||||
if (name.contains(".")) { // Nur wenn der Linktext wie ein Dateiname aussieht
|
||||
String type = hint != null ? inferTypeFromKey(hint) : null;
|
||||
|
||||
// NEW: Falls Typ unklar, versuche ihn aus der URL zu extrahieren (speziell HF)
|
||||
if ((type == null || type.equals("checkpoints")) && url.contains("/resolve/main/split_files/")) {
|
||||
String sub = url.substring(url.indexOf("/split_files/") + 13);
|
||||
if (sub.contains("/")) type = sub.substring(0, sub.indexOf("/"));
|
||||
} else if ((type == null || type.equals("checkpoints")) && url.contains("/resolve/main/")) {
|
||||
// Generischer Check für HF Ordner-Strukturen
|
||||
String sub = url.substring(url.indexOf("/resolve/main/") + 14);
|
||||
if (sub.contains("/")) {
|
||||
String folder = sub.substring(0, sub.indexOf("/")).toLowerCase();
|
||||
if (folder.equals("loras") || folder.equals("vae") || folder.equals("checkpoints") || folder.equals("text_encoders") || folder.equals("diffusion_models") || folder.equals("controlnet")) {
|
||||
type = folder;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (type == null) type = ctx;
|
||||
addModelInfo(res, new ModelInfo(type != null ? type : "checkpoints", name, url));
|
||||
}
|
||||
}
|
||||
|
||||
Matcher m = ExpertPatterns.FILE.matcher(v);
|
||||
while (m.find()) {
|
||||
String f = m.group(1);
|
||||
String match = m.group(1);
|
||||
|
||||
// Ignoriere offensichtliche URLs (wir wollen nur Dateinamen)
|
||||
if (match.startsWith("http") || match.startsWith("//")) continue;
|
||||
|
||||
// Extrahiere nur den reinen Dateinamen, falls ein Pfad gematcht wurde
|
||||
String f = match;
|
||||
if (f.contains("/")) f = f.substring(f.lastIndexOf("/") + 1);
|
||||
if (f.contains("\\")) f = f.substring(f.lastIndexOf("\\") + 1);
|
||||
|
||||
if (f.length() > 3) {
|
||||
// Erneuter Check nach Pfad-Extraktion
|
||||
String fLow = f.toLowerCase();
|
||||
if (fLow.equals("none.safetensors") || fLow.equals("default.safetensors")) continue;
|
||||
|
||||
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";
|
||||
|
||||
// Verfeinerung basierend auf Dateinamen, aber nur wenn der Typ noch unklar oder zu allgemein ist
|
||||
if (fLow.contains("vae_approx") || fLow.contains("vae-approx")) type = "vae_approx";
|
||||
else if (ExpertPatterns.COMP_VAE.matcher(fLow).find() && (type == null || type.equals("checkpoints"))) type = "vae";
|
||||
else if (fLow.contains("clip_vision")) type = "clip_vision";
|
||||
else if (ExpertPatterns.COMP_CLIP.matcher(fLow).find() && (type == null || type.equals("checkpoints"))) type = "clip";
|
||||
else if (ExpertPatterns.COMP_UNET.matcher(fLow).find() && (type == null || type.equals("checkpoints") || (type.equals("clip") && !fLow.contains("clip_l") && !fLow.contains("clip_g")))) type = "unet";
|
||||
|
||||
addModelInfo(res, new ModelInfo(type != null ? type : "checkpoints", f, "MISSING"));
|
||||
}
|
||||
}
|
||||
@@ -160,30 +271,63 @@ public class ComfyModelAnalyzer implements IModelAnalyzer {
|
||||
private String inferTypeFromNode(String nt) {
|
||||
nt = nt.toLowerCase();
|
||||
if (nt.contains("lora")) return "loras";
|
||||
if (nt.contains("vae_approx") || nt.contains("vae-approx")) return "vae_approx";
|
||||
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";
|
||||
if (nt.contains("latent_upscale")) return "latent_upscale_models";
|
||||
if (nt.contains("clip_vision")) return "clip_vision";
|
||||
if (nt.contains("clip") || nt.contains("text_encoder") || nt.contains("cliploader") || nt.contains("llama") || nt.contains("mistral") || nt.contains("t5") || nt.contains("qwen")) return "clip";
|
||||
if (nt.contains("diffusion_model")) return "diffusion_models";
|
||||
if (nt.contains("unet") || nt.contains("diffusion")) return "unet";
|
||||
if (nt.contains("gligen")) return "gligen";
|
||||
if (nt.contains("embeddings") || nt.contains("textual_inversion")) return "embeddings";
|
||||
if (nt.contains("hypernetwork")) return "hypernetworks";
|
||||
if (nt.contains("photomaker")) return "photomaker";
|
||||
if (nt.contains("style_model")) return "style_models";
|
||||
if (nt.contains("audio") && nt.contains("encoder")) return "audio_encoders";
|
||||
if (nt.contains("diffusers")) return "diffusers";
|
||||
if (nt.contains("configs")) return "configs";
|
||||
if (nt.contains("model_patch")) return "model_patches";
|
||||
return null;
|
||||
}
|
||||
|
||||
private String inferTypeFromKey(String k) {
|
||||
k = k.toLowerCase();
|
||||
if (k.contains("lora")) return "loras";
|
||||
if (k.contains("vae_approx") || k.contains("vae-approx")) return "vae_approx";
|
||||
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("clip_vision")) return "clip_vision";
|
||||
if (k.contains("clip") || k.contains("text_encoder") || k.contains("llama") || k.contains("mistral") || k.contains("t5") || k.contains("qwen")) return "clip";
|
||||
if (k.contains("diffusion_model")) return "diffusion_models";
|
||||
if (k.contains("unet") || k.contains("diffusion") || k.contains("hidream") || k.contains("aura") || k.contains("mochi") || k.contains("svd") || k.contains("z_image") || k.contains("lumina") || k.contains("flux") || k.contains("dit")) return "unet";
|
||||
if (k.contains("control") && k.contains("net")) return "controlnet";
|
||||
if (k.contains("upscale") || k.contains("esrgan")) return "upscale_models";
|
||||
if (k.contains("embedding")) return "embeddings";
|
||||
if (k.contains("hypernetwork")) return "hypernetworks";
|
||||
if (k.contains("gligen")) return "gligen";
|
||||
if (k.contains("photomaker")) return "photomaker";
|
||||
if (k.contains("style_model")) return "style_models";
|
||||
if (k.contains("audio") && k.contains("encoder")) return "audio_encoders";
|
||||
if (k.contains("config")) return "configs";
|
||||
return null;
|
||||
}
|
||||
|
||||
private void addModelInfo(List<ModelInfo> res, ModelInfo info) {
|
||||
for (ModelInfo e : res) {
|
||||
if (e.getName().equalsIgnoreCase(info.getName())) {
|
||||
if (e.getUrl().equals("MISSING") && !info.getUrl().equals("MISSING")) e.setUrl(info.getUrl());
|
||||
if (e.getType().equals("checkpoints") && !info.getType().equals("checkpoints")) e.setType(info.getType());
|
||||
// Wenn die neue Info eine echte URL hat, ist sie wahrscheinlich aus Metadaten -> Vorrang!
|
||||
if (!info.getUrl().equals("MISSING")) {
|
||||
e.setUrl(info.getUrl());
|
||||
e.setType(info.getType());
|
||||
e.setSave_path(info.getSave_path());
|
||||
} else if (e.getUrl().equals("MISSING")) {
|
||||
// Falls beide keine URL haben, nimm den spezifischeren Pfad (nicht checkpoints)
|
||||
if ((e.getType() == null || e.getType().equals("checkpoints")) && info.getType() != null && !info.getType().equals("checkpoints")) {
|
||||
e.setType(info.getType());
|
||||
e.setSave_path(info.getSave_path());
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,35 +4,74 @@ import org.json.JSONObject;
|
||||
import org.springframework.stereotype.Service;
|
||||
|
||||
import java.io.File;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
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;
|
||||
private static final String VAULT_FILE = "settings.vault";
|
||||
private JSONObject settings = new JSONObject();
|
||||
private String masterPassword = null;
|
||||
|
||||
public ConfigService() {
|
||||
load();
|
||||
/**
|
||||
* Use the current working directory for application data.
|
||||
*/
|
||||
public String getAppDataPath() {
|
||||
return System.getProperty("user.dir");
|
||||
}
|
||||
|
||||
private void load() {
|
||||
try {
|
||||
File file = new File(CONFIG_FILE);
|
||||
if (file.exists()) {
|
||||
settings = new JSONObject(Files.readString(file.toPath()));
|
||||
public File getFileInAppData(String filename) {
|
||||
return new File(getAppDataPath(), filename);
|
||||
}
|
||||
|
||||
public void unlock(String password) throws Exception {
|
||||
File vault = getFileInAppData(VAULT_FILE);
|
||||
System.out.println("Attempting to unlock vault at: " + vault.getAbsolutePath());
|
||||
|
||||
if (vault.exists()) {
|
||||
String encrypted = Files.readString(vault.toPath(), StandardCharsets.UTF_8);
|
||||
try {
|
||||
String decrypted = EncryptionUtils.decrypt(encrypted, password);
|
||||
JSONObject decryptedJson = new JSONObject(decrypted);
|
||||
this.settings = decryptedJson;
|
||||
this.masterPassword = password;
|
||||
System.out.println("Vault unlocked successfully. Keys found: " + settings.keySet());
|
||||
} catch (Exception e) {
|
||||
System.err.println("Failed to unlock vault: " + e.getMessage());
|
||||
throw new Exception("Wrong password or corrupted vault!");
|
||||
}
|
||||
} else {
|
||||
this.masterPassword = password;
|
||||
File oldFile = getFileInAppData(CONFIG_FILE);
|
||||
if (oldFile.exists()) {
|
||||
try {
|
||||
String content = Files.readString(oldFile.toPath(), StandardCharsets.UTF_8);
|
||||
this.settings = new JSONObject(content);
|
||||
save();
|
||||
oldFile.delete();
|
||||
System.out.println("Migrated old settings to encrypted vault.");
|
||||
} catch (Exception e) {
|
||||
System.err.println("Failed to migrate old settings: " + e.getMessage());
|
||||
}
|
||||
} else {
|
||||
settings = new JSONObject();
|
||||
System.out.println("No vault found, initialized new empty vault at: " + vault.getAbsolutePath());
|
||||
save();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
settings = new JSONObject();
|
||||
}
|
||||
}
|
||||
|
||||
public void save() {
|
||||
if (masterPassword == null) {
|
||||
System.err.println("Cannot save: Vault not unlocked.");
|
||||
return;
|
||||
}
|
||||
try {
|
||||
Files.writeString(Paths.get(CONFIG_FILE), settings.toString(4));
|
||||
String encrypted = EncryptionUtils.encrypt(settings.toString(), masterPassword);
|
||||
Files.writeString(getFileInAppData(VAULT_FILE).toPath(), encrypted, StandardCharsets.UTF_8);
|
||||
System.out.println("Vault saved successfully to: " + getFileInAppData(VAULT_FILE).getAbsolutePath());
|
||||
} catch (Exception e) {
|
||||
System.err.println("Error saving vault: " + e.getMessage());
|
||||
e.printStackTrace();
|
||||
}
|
||||
}
|
||||
@@ -42,4 +81,11 @@ public class ConfigService {
|
||||
|
||||
public String getHfToken() { return settings.optString("hf_token", ""); }
|
||||
public void setHfToken(String token) { settings.put("hf_token", token); save(); }
|
||||
}
|
||||
|
||||
public String getModelsPath() { return settings.optString("models_path", "C:\\pinokio\\api\\comfy.git\\app\\models"); }
|
||||
public void setModelsPath(String path) { settings.put("models_path", path); save(); }
|
||||
|
||||
public boolean isUnlocked() {
|
||||
return masterPassword != null;
|
||||
}
|
||||
}
|
||||
@@ -22,28 +22,78 @@ public class DefaultDownloadManager implements IDownloadManager {
|
||||
private final HttpClient httpClient = HttpClient.newBuilder().followRedirects(HttpClient.Redirect.ALWAYS).build();
|
||||
private volatile boolean isPaused = false;
|
||||
private volatile boolean isStopped = false;
|
||||
private volatile boolean[] currentSelection;
|
||||
private final java.util.Set<Integer> completedIndices = java.util.concurrent.ConcurrentHashMap.newKeySet();
|
||||
|
||||
@Autowired
|
||||
private ConfigService configService;
|
||||
|
||||
private void safeUpdateStatus(int index, String status, BiConsumer<Integer, String> statusUpdater) {
|
||||
if (completedIndices.contains(index)) return;
|
||||
|
||||
if (status.contains("✅ Finished") || status.contains("Stopped") || status.contains("Skipped") || status.contains("Already exists")) {
|
||||
completedIndices.add(index);
|
||||
}
|
||||
statusUpdater.accept(index, status);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void startQueue(List<ModelInfo> models, String baseDir, BiConsumer<Integer, String> statusUpdater, Runnable onFinished) {
|
||||
public void startQueue(List<ModelInfo> models, boolean[] selectedIndices, String baseDir, BiConsumer<Integer, String> statusUpdater, Runnable onFinished) {
|
||||
isStopped = false;
|
||||
isPaused = false;
|
||||
this.currentSelection = selectedIndices;
|
||||
completedIndices.clear();
|
||||
new Thread(() -> {
|
||||
try {
|
||||
java.util.concurrent.CompletableFuture<?>[] futures = new java.util.concurrent.CompletableFuture[models.size()];
|
||||
for (int i = 0; i < models.size(); i++) {
|
||||
if (isStopped) {
|
||||
for (int j = i; j < models.size(); j++) {
|
||||
futures[j] = java.util.concurrent.CompletableFuture.completedFuture(null);
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
while (isPaused && !isStopped) {
|
||||
try {
|
||||
Thread.sleep(200);
|
||||
} catch (InterruptedException e) {
|
||||
Thread.currentThread().interrupt();
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
final int index = i;
|
||||
|
||||
if (!isSelected(index)) {
|
||||
safeUpdateStatus(index, "Skipped (Not Selected)", statusUpdater);
|
||||
futures[i] = java.util.concurrent.CompletableFuture.completedFuture(null);
|
||||
continue;
|
||||
}
|
||||
|
||||
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; }
|
||||
if (isStopped) { safeUpdateStatus(index, "Stopped", statusUpdater); return; }
|
||||
if (!isSelected(index)) { safeUpdateStatus(index, "Skipped (Not Selected)", statusUpdater); return; }
|
||||
if (info.getUrl() == null || info.getUrl().startsWith("MISSING")) { safeUpdateStatus(index, "Skipped (No URL)", statusUpdater); return; }
|
||||
try {
|
||||
Path targetDir = Paths.get(baseDir, info.getType());
|
||||
String subPath = info.getSave_path();
|
||||
if (subPath == null || subPath.isEmpty() || "default".equalsIgnoreCase(subPath)) {
|
||||
subPath = info.getType() != null ? info.getType() : "checkpoints";
|
||||
}
|
||||
|
||||
String finalBaseDir = baseDir != null ? baseDir : ".";
|
||||
Path targetDir = Paths.get(finalBaseDir, subPath);
|
||||
Files.createDirectories(targetDir);
|
||||
downloadWithResume(info, targetDir.resolve(info.getName()), index, statusUpdater);
|
||||
} catch (Exception e) { statusUpdater.accept(index, "Error: " + e.getMessage()); }
|
||||
downloadWithResume(info, targetDir.resolve(info.getName()), index, statusUpdater);
|
||||
} catch (Exception e) {
|
||||
if (isStopped || !isSelected(index)) {
|
||||
safeUpdateStatus(index, !isSelected(index) ? "Skipped (Unchecked)" : "Stopped", statusUpdater);
|
||||
} else {
|
||||
String msg = e.getMessage() != null ? e.getMessage() : e.getClass().getSimpleName();
|
||||
safeUpdateStatus(index, "Error: " + msg, statusUpdater);
|
||||
}
|
||||
}
|
||||
}, executor);
|
||||
}
|
||||
java.util.concurrent.CompletableFuture.allOf(futures).join();
|
||||
@@ -51,62 +101,155 @@ public class DefaultDownloadManager implements IDownloadManager {
|
||||
}).start();
|
||||
}
|
||||
|
||||
private void downloadWithResume(ModelInfo info, Path targetFile, int index, BiConsumer<Integer, String> statusUpdater) throws Exception {
|
||||
File file = targetFile.toFile();
|
||||
|
||||
// ANTI-STUB LOGIK: Wenn Datei < 10KB, löschen wir sie und fangen neu an (da es wahrscheinlich ein LFS-Pointer ist)
|
||||
if (file.exists() && file.length() < 10240) {
|
||||
file.delete();
|
||||
}
|
||||
|
||||
long existingFileSize = file.exists() ? file.length() : 0;
|
||||
|
||||
HttpRequest.Builder builder = HttpRequest.newBuilder().uri(URI.create(info.getUrl())).header("User-Agent", "Mozilla/5.0");
|
||||
String hfToken = configService.getHfToken();
|
||||
if (info.getUrl().contains("huggingface.co") && !hfToken.isEmpty()) {
|
||||
builder.header("Authorization", "Bearer " + hfToken);
|
||||
}
|
||||
private boolean isSelected(int index) {
|
||||
return currentSelection == null || (index < currentSelection.length && currentSelection[index]);
|
||||
}
|
||||
|
||||
if (existingFileSize > 0) builder.header("Range", "bytes=" + existingFileSize + "-");
|
||||
@Override
|
||||
public void updateSelection(boolean[] selectedIndices) {
|
||||
this.currentSelection = selectedIndices;
|
||||
}
|
||||
|
||||
HttpResponse<InputStream> response = httpClient.send(builder.build(), HttpResponse.BodyHandlers.ofInputStream());
|
||||
int statusCode = response.statusCode();
|
||||
|
||||
if (statusCode == 401 || statusCode == 403) {
|
||||
statusUpdater.accept(index, "❌ Auth Required (Token?)");
|
||||
return;
|
||||
}
|
||||
private void downloadWithResume(ModelInfo info, Path targetFile, int index, BiConsumer<Integer, String> statusUpdater) {
|
||||
try {
|
||||
if (waitForPauseAndCheckSelection(index, statusUpdater)) return;
|
||||
|
||||
long contentLen = response.headers().firstValueAsLong("Content-Length").orElse(0L);
|
||||
long totalBytes = (statusCode == 206) ? contentLen + existingFileSize : contentLen;
|
||||
File file = targetFile.toFile();
|
||||
String hfToken = configService.getHfToken();
|
||||
|
||||
// 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 (file.exists() && file.length() < 10240) {
|
||||
file.delete();
|
||||
}
|
||||
|
||||
if (statusCode == 200) existingFileSize = 0;
|
||||
long existingFileSize = file.exists() ? file.length() : 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;
|
||||
if (waitForPauseAndCheckSelection(index, statusUpdater)) return;
|
||||
|
||||
HttpRequest.Builder headBuilder = HttpRequest.newBuilder()
|
||||
.uri(URI.create(info.getUrl()))
|
||||
.header("User-Agent", "Mozilla/5.0")
|
||||
.method("HEAD", HttpRequest.BodyPublishers.noBody());
|
||||
|
||||
if (info.getUrl().contains("huggingface.co") && hfToken != null && !hfToken.isEmpty()) {
|
||||
headBuilder.header("Authorization", "Bearer " + hfToken);
|
||||
}
|
||||
|
||||
HttpResponse<Void> headResponse = httpClient.send(headBuilder.build(), HttpResponse.BodyHandlers.discarding());
|
||||
|
||||
if (headResponse.statusCode() == 401 || headResponse.statusCode() == 403) {
|
||||
safeUpdateStatus(index, "❌ Auth Required (Token?)", statusUpdater);
|
||||
return;
|
||||
}
|
||||
|
||||
long totalRemoteSize = headResponse.headers().firstValueAsLong("Content-Length").orElse(0L);
|
||||
|
||||
if (existingFileSize > 0 && totalRemoteSize > 0 && existingFileSize == totalRemoteSize) {
|
||||
safeUpdateStatus(index, "✅ Already exists", statusUpdater);
|
||||
return;
|
||||
}
|
||||
|
||||
if (waitForPauseAndCheckSelection(index, statusUpdater)) return;
|
||||
|
||||
HttpRequest.Builder downloadBuilder = HttpRequest.newBuilder()
|
||||
.uri(URI.create(info.getUrl()))
|
||||
.header("User-Agent", "Mozilla/5.0");
|
||||
|
||||
if (info.getUrl().contains("huggingface.co") && hfToken != null && !hfToken.isEmpty()) {
|
||||
downloadBuilder.header("Authorization", "Bearer " + hfToken);
|
||||
}
|
||||
|
||||
if (existingFileSize > 0) downloadBuilder.header("Range", "bytes=" + existingFileSize + "-");
|
||||
|
||||
HttpResponse<InputStream> response = httpClient.send(downloadBuilder.build(), HttpResponse.BodyHandlers.ofInputStream());
|
||||
int statusCode = response.statusCode();
|
||||
|
||||
if (statusCode == 416) {
|
||||
safeUpdateStatus(index, "✅ Already exists", statusUpdater);
|
||||
return;
|
||||
}
|
||||
|
||||
if (statusCode == 401 || statusCode == 403) {
|
||||
safeUpdateStatus(index, "❌ Auth Required (Token?)", statusUpdater);
|
||||
return;
|
||||
}
|
||||
|
||||
long contentLen = response.headers().firstValueAsLong("Content-Length").orElse(0L);
|
||||
long totalBytes = (statusCode == 206) ? contentLen + existingFileSize : contentLen;
|
||||
|
||||
if (totalBytes > 0 && totalBytes < 5000 && info.getName().endsWith(".safetensors")) {
|
||||
safeUpdateStatus(index, "❌ LFS Stub detected (Token?)", statusUpdater);
|
||||
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 || !isSelected(index) || Thread.currentThread().isInterrupted()) {
|
||||
safeUpdateStatus(index, !isSelected(index) ? "Skipped (Unchecked)" : "Stopped", statusUpdater);
|
||||
return;
|
||||
}
|
||||
|
||||
if (isPaused) {
|
||||
if (waitForPauseAndCheckSelection(index, statusUpdater)) return;
|
||||
safeUpdateStatus(index, "Resuming...", statusUpdater);
|
||||
}
|
||||
|
||||
raf.write(buffer, 0, read);
|
||||
downloaded += read;
|
||||
|
||||
long now = System.currentTimeMillis();
|
||||
if (now - lastUpdate > 800) {
|
||||
safeUpdateStatus(index, "Downloading: " + (totalBytes > 0 ? (downloaded * 100 / totalBytes) : "?") + "% (" + formatSize(downloaded) + ")", statusUpdater);
|
||||
lastUpdate = now;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (isStopped || !isSelected(index)) {
|
||||
safeUpdateStatus(index, !isSelected(index) ? "Skipped (Unchecked)" : "Stopped", statusUpdater);
|
||||
} else {
|
||||
if (totalBytes > 0 && file.length() < totalBytes) {
|
||||
safeUpdateStatus(index, "❌ Incomplete (" + formatSize(file.length()) + "/" + formatSize(totalBytes) + ")", statusUpdater);
|
||||
} else {
|
||||
safeUpdateStatus(index, "✅ Finished", statusUpdater);
|
||||
}
|
||||
}
|
||||
} catch (Exception e) {
|
||||
// Falls es ein Timing-Problem ist (Flagge wird gesetzt während Exception fliegt)
|
||||
try { Thread.sleep(50); } catch (InterruptedException ignored) {}
|
||||
|
||||
if (isStopped || !isSelected(index) || Thread.currentThread().isInterrupted()) {
|
||||
safeUpdateStatus(index, !isSelected(index) ? "Skipped (Unchecked)" : "Stopped", statusUpdater);
|
||||
} else {
|
||||
String msg = e.getMessage() != null ? e.getMessage() : e.getClass().getSimpleName();
|
||||
safeUpdateStatus(index, "Error: " + msg, statusUpdater);
|
||||
}
|
||||
}
|
||||
statusUpdater.accept(index, "Finished");
|
||||
}
|
||||
|
||||
private boolean waitForPauseAndCheckSelection(int index, BiConsumer<Integer, String> statusUpdater) {
|
||||
try {
|
||||
if (isPaused) {
|
||||
safeUpdateStatus(index, "Paused", statusUpdater);
|
||||
}
|
||||
while (isPaused && !isStopped && isSelected(index)) {
|
||||
Thread.sleep(200);
|
||||
}
|
||||
} catch (InterruptedException e) {
|
||||
Thread.currentThread().interrupt();
|
||||
}
|
||||
|
||||
if (isStopped || !isSelected(index)) {
|
||||
safeUpdateStatus(index, !isSelected(index) ? "Skipped (Unchecked)" : "Stopped", statusUpdater);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
private String formatSize(long bytes) {
|
||||
@@ -118,4 +261,4 @@ public class DefaultDownloadManager implements IDownloadManager {
|
||||
@Override public void togglePause() { isPaused = !isPaused; }
|
||||
@Override public void stop() { isStopped = true; isPaused = false; }
|
||||
@Override public boolean isPaused() { return isPaused; }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package de.tki.comfymodels.service.impl;
|
||||
|
||||
import javax.crypto.Cipher;
|
||||
import javax.crypto.SecretKey;
|
||||
import javax.crypto.SecretKeyFactory;
|
||||
import javax.crypto.spec.PBEKeySpec;
|
||||
import javax.crypto.spec.SecretKeySpec;
|
||||
import javax.crypto.spec.IvParameterSpec;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.security.SecureRandom;
|
||||
import java.security.spec.KeySpec;
|
||||
import java.util.Base64;
|
||||
|
||||
public class EncryptionUtils {
|
||||
private static final int ITERATIONS = 65536;
|
||||
private static final int KEY_LENGTH = 256;
|
||||
private static final byte[] SALT = "ComfyUI-Vault-Salt-2026".getBytes(StandardCharsets.UTF_8);
|
||||
|
||||
public static String encrypt(String strToEncrypt, String password) throws Exception {
|
||||
byte[] iv = new byte[16];
|
||||
new SecureRandom().nextBytes(iv);
|
||||
IvParameterSpec ivspec = new IvParameterSpec(iv);
|
||||
|
||||
SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
|
||||
KeySpec spec = new PBEKeySpec(password.toCharArray(), SALT, ITERATIONS, KEY_LENGTH);
|
||||
SecretKey tmp = factory.generateSecret(spec);
|
||||
SecretKeySpec secretKey = new SecretKeySpec(tmp.getEncoded(), "AES");
|
||||
|
||||
Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
|
||||
cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivspec);
|
||||
|
||||
byte[] encrypted = cipher.doFinal(strToEncrypt.getBytes(StandardCharsets.UTF_8));
|
||||
byte[] combined = new byte[iv.length + encrypted.length];
|
||||
System.arraycopy(iv, 0, combined, 0, iv.length);
|
||||
System.arraycopy(encrypted, 0, combined, iv.length, encrypted.length);
|
||||
|
||||
return Base64.getEncoder().encodeToString(combined);
|
||||
}
|
||||
|
||||
public static String decrypt(String strToDecrypt, String password) throws Exception {
|
||||
if (strToDecrypt == null || strToDecrypt.trim().isEmpty()) {
|
||||
throw new Exception("Encrypted data is empty");
|
||||
}
|
||||
byte[] combined = Base64.getDecoder().decode(strToDecrypt.trim());
|
||||
if (combined.length < 16) {
|
||||
throw new Exception("Encrypted data too short");
|
||||
}
|
||||
|
||||
byte[] iv = new byte[16];
|
||||
System.arraycopy(combined, 0, iv, 0, iv.length);
|
||||
IvParameterSpec ivspec = new IvParameterSpec(iv);
|
||||
|
||||
byte[] encrypted = new byte[combined.length - 16];
|
||||
System.arraycopy(combined, 16, encrypted, 0, encrypted.length);
|
||||
|
||||
SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
|
||||
KeySpec spec = new PBEKeySpec(password.toCharArray(), SALT, ITERATIONS, KEY_LENGTH);
|
||||
SecretKey tmp = factory.generateSecret(spec);
|
||||
SecretKeySpec secretKey = new SecretKeySpec(tmp.getEncoded(), "AES");
|
||||
|
||||
Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
|
||||
cipher.init(Cipher.DECRYPT_MODE, secretKey, ivspec);
|
||||
return new String(cipher.doFinal(encrypted), StandardCharsets.UTF_8);
|
||||
}
|
||||
}
|
||||
@@ -62,7 +62,8 @@ public class GeminiAIService {
|
||||
if (apiKey == null || apiKey.trim().isEmpty()) return null;
|
||||
|
||||
try {
|
||||
String shortContext = metadataContext.length() > 25000 ? metadataContext.substring(0, 25000) : metadataContext;
|
||||
String context = metadataContext != null ? metadataContext : "No context";
|
||||
String shortContext = context.length() > 25000 ? context.substring(0, 25000) : context;
|
||||
|
||||
String prompt = "Websuche: Ermittle das offizielle Hugging Face Repository für die Datei '" + modelName + "'.\n\n" +
|
||||
"KONTEXT:\n" +
|
||||
|
||||
@@ -3,6 +3,7 @@ package de.tki.comfymodels.service.impl;
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import org.json.JSONArray;
|
||||
import org.json.JSONObject;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.stereotype.Service;
|
||||
|
||||
import jakarta.annotation.PostConstruct;
|
||||
@@ -18,6 +19,9 @@ public class ModelListService {
|
||||
private static final String STORAGE_FILE = "uploaded_models.json";
|
||||
private List<ModelInfo> models = new ArrayList<>();
|
||||
|
||||
@Autowired
|
||||
private ConfigService configService;
|
||||
|
||||
@PostConstruct
|
||||
public void init() {
|
||||
loadFromStorage();
|
||||
@@ -31,13 +35,14 @@ public class ModelListService {
|
||||
List<ModelInfo> newModels = new ArrayList<>();
|
||||
for (int i = 0; i < modelsArray.length(); i++) {
|
||||
JSONObject m = modelsArray.getJSONObject(i);
|
||||
String rawDir = m.optString("directory", m.optString("type", "checkpoints"));
|
||||
ModelInfo info = new ModelInfo(
|
||||
m.optString("type", "Unknown"),
|
||||
rawDir,
|
||||
m.optString("name", "Unknown"),
|
||||
m.optString("url", "MISSING")
|
||||
);
|
||||
info.setBase(m.optString("base", ""));
|
||||
info.setSave_path(m.optString("save_path", ""));
|
||||
info.setSave_path(m.optString("save_path", rawDir));
|
||||
info.setDescription(m.optString("description", ""));
|
||||
info.setReference(m.optString("reference", ""));
|
||||
info.setFilename(m.optString("filename", ""));
|
||||
@@ -57,16 +62,17 @@ public class ModelListService {
|
||||
|
||||
private void loadFromStorage() {
|
||||
try {
|
||||
File file = new File(STORAGE_FILE);
|
||||
File file = configService != null ? configService.getFileInAppData(STORAGE_FILE) : new File(STORAGE_FILE);
|
||||
if (file.exists()) {
|
||||
String content = Files.readString(file.toPath(), StandardCharsets.UTF_8);
|
||||
JSONArray array = new JSONArray(content);
|
||||
models.clear();
|
||||
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")
|
||||
m.optString("type", "checkpoints"),
|
||||
m.optString("name", "Unknown"),
|
||||
m.optString("url", "MISSING")
|
||||
);
|
||||
info.setBase(m.optString("base"));
|
||||
info.setSave_path(m.optString("save_path"));
|
||||
@@ -98,7 +104,8 @@ public class ModelListService {
|
||||
jo.put("size", m.getSize());
|
||||
array.put(jo);
|
||||
}
|
||||
Files.writeString(new File(STORAGE_FILE).toPath(), array.toString(4), StandardCharsets.UTF_8);
|
||||
File file = configService != null ? configService.getFileInAppData(STORAGE_FILE) : new File(STORAGE_FILE);
|
||||
Files.writeString(file.toPath(), array.toString(4), StandardCharsets.UTF_8);
|
||||
} catch (Exception e) {
|
||||
e.printStackTrace();
|
||||
}
|
||||
|
||||
@@ -14,6 +14,9 @@ import java.net.http.HttpRequest;
|
||||
import java.net.http.HttpResponse;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.*;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.function.BiConsumer;
|
||||
|
||||
@Service
|
||||
@@ -39,80 +42,133 @@ public class ModelSearchService implements IModelSearchService {
|
||||
.followRedirects(HttpClient.Redirect.ALWAYS)
|
||||
.build();
|
||||
|
||||
private final ExecutorService searchExecutor = Executors.newFixedThreadPool(4);
|
||||
|
||||
@Override
|
||||
public void searchOnline(List<ModelInfo> modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName,
|
||||
BiConsumer<Integer, String> onStatusUpdate,
|
||||
BiConsumer<Integer, ModelInfo> onModelFound) {
|
||||
if (modelsToDownload == null) return;
|
||||
|
||||
for (int i = 0; i < modelsToDownload.size(); i++) {
|
||||
if (!selectedIndices[i]) continue;
|
||||
|
||||
ModelInfo info = modelsToDownload.get(i);
|
||||
final int index = i;
|
||||
new Thread(() -> {
|
||||
// Priority 1: User defined Model List
|
||||
Optional<ModelInfo> manualMatch = modelListService.findByFilename(info.getName());
|
||||
if (manualMatch.isPresent() && !manualMatch.get().getUrl().equals("MISSING")) {
|
||||
onStatusUpdate.accept(index, "📂 Found in Model List");
|
||||
if (validateAndSetUrl(info, index, manualMatch.get().getUrl(), "📂 USER DEFINED", onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
onStatusUpdate.accept(index, "✨ Gemini Scouting...");
|
||||
String aiHint = geminiService.discoverBestRepo(info.getName(), fileName, workflowContext);
|
||||
if (aiHint != null && !aiHint.equalsIgnoreCase("UNKNOWN")) {
|
||||
if (aiHint.startsWith("http")) {
|
||||
if (validateAndSetUrl(info, index, aiHint, "✨ AI DIRECT", onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
onStatusUpdate.accept(index, "🔍 Validating Repo: " + aiHint);
|
||||
if (fetchHuggingFaceUrlInSpecificRepo(info, index, aiHint, onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
String modelName = info.getName();
|
||||
String cleanName = modelName.replaceAll("(_fp8|_fp16|_bf16|_v\\d+|\\d+v|_fix|\\.safetensors|\\.sft|\\.ckpt)", "");
|
||||
String baseName = cleanName.split("[_\\- ]")[0];
|
||||
|
||||
String contextPrefix = fileName.toLowerCase().split("[_\\- ]")[0];
|
||||
if (contextPrefix.length() < 3) contextPrefix = "";
|
||||
|
||||
LinkedHashSet<String> attempts = new LinkedHashSet<>();
|
||||
if (!contextPrefix.isEmpty()) {
|
||||
attempts.add(contextPrefix + " " + modelName);
|
||||
attempts.add(contextPrefix + " " + cleanName);
|
||||
}
|
||||
attempts.add(modelName);
|
||||
attempts.add(cleanName);
|
||||
attempts.add(baseName);
|
||||
|
||||
// Suche in offiziellen Repos (Top 100)
|
||||
for (String query : attempts) {
|
||||
if (query.length() < 3) continue;
|
||||
onStatusUpdate.accept(index, "🔍 Searching Official: " + query);
|
||||
if (fetchHuggingFaceUrl(info, index, query, true, onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
// Suche in Community Repos (Top 100)
|
||||
for (String query : attempts) {
|
||||
if (query.length() < 3) continue;
|
||||
onStatusUpdate.accept(index, "🔍 Searching Community: " + query);
|
||||
if (fetchHuggingFaceUrl(info, index, query, false, onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
// Civitai Fallback
|
||||
for (String query : attempts) {
|
||||
if (query.length() < 3) continue;
|
||||
onStatusUpdate.accept(index, "🔍 Searching Civitai: " + query);
|
||||
if (fetchCivitaiUrl(info, index, query, onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
onStatusUpdate.accept(index, "❌ No trusted match found");
|
||||
}).start();
|
||||
BiConsumer<Integer, String> onStatusUpdate,
|
||||
BiConsumer<Integer, ModelInfo> onModelFound,
|
||||
Runnable onFinished) {
|
||||
if (modelsToDownload == null) {
|
||||
if (onFinished != null) onFinished.run();
|
||||
return;
|
||||
}
|
||||
|
||||
List<Integer> targetIndices = new ArrayList<>();
|
||||
for (int i = 0; i < modelsToDownload.size(); i++) {
|
||||
if (selectedIndices != null && i < selectedIndices.length && !selectedIndices[i]) continue;
|
||||
targetIndices.add(i);
|
||||
}
|
||||
|
||||
if (targetIndices.isEmpty()) {
|
||||
if (onFinished != null) onFinished.run();
|
||||
return;
|
||||
}
|
||||
|
||||
AtomicInteger remaining = new AtomicInteger(targetIndices.size());
|
||||
|
||||
for (int index : targetIndices) {
|
||||
ModelInfo info = modelsToDownload.get(index);
|
||||
searchExecutor.submit(() -> {
|
||||
try {
|
||||
performSearch(info, index, fileName, workflowContext, onStatusUpdate, onModelFound);
|
||||
} catch (Exception e) {
|
||||
onStatusUpdate.accept(index, "Error: " + e.getMessage());
|
||||
} finally {
|
||||
if (remaining.decrementAndGet() == 0 && onFinished != null) {
|
||||
onFinished.run();
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private void performSearch(ModelInfo info, int index, String fileName, String workflowContext,
|
||||
BiConsumer<Integer, String> onStatusUpdate,
|
||||
BiConsumer<Integer, ModelInfo> onModelFound) {
|
||||
// Priority 1: User defined Model List
|
||||
Optional<ModelInfo> manualMatch = modelListService.findByFilename(info.getName());
|
||||
if (manualMatch.isPresent() && !manualMatch.get().getUrl().equals("MISSING")) {
|
||||
onStatusUpdate.accept(index, "📂 Found in Model List");
|
||||
if (validateAndSetUrl(info, index, manualMatch.get().getUrl(), "📂 USER DEFINED", onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
onStatusUpdate.accept(index, "✨ Gemini Scouting...");
|
||||
String aiHint = geminiService.discoverBestRepo(info.getName(), fileName, workflowContext);
|
||||
if (aiHint != null && !aiHint.equalsIgnoreCase("UNKNOWN")) {
|
||||
if (aiHint.startsWith("http")) {
|
||||
if (validateAndSetUrl(info, index, aiHint, "✨ AI DIRECT", onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
onStatusUpdate.accept(index, "🔍 Validating Repo: " + aiHint);
|
||||
if (fetchHuggingFaceUrlInSpecificRepo(info, index, aiHint, onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
String modelName = info.getName();
|
||||
String cleanName = modelName.replaceAll("(_fp8|_fp16|_bf16|_v\\d+|\\d+v|_fix|\\.safetensors|\\.sft|\\.ckpt)", "");
|
||||
String[] nameParts = cleanName.split("[_\\- ]");
|
||||
String baseName = nameParts.length > 0 ? nameParts[0] : cleanName;
|
||||
|
||||
String contextPrefix = "";
|
||||
if (fileName != null && !fileName.isEmpty()) {
|
||||
String[] fileParts = fileName.toLowerCase().split("[_\\- ]");
|
||||
if (fileParts.length > 0) {
|
||||
contextPrefix = fileParts[0];
|
||||
if (contextPrefix.length() < 3 || contextPrefix.startsWith("workflow")) {
|
||||
contextPrefix = "";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (contextPrefix.isEmpty() && info.getPopularity() != null && info.getPopularity().contains("Context: ")) {
|
||||
String[] popParts = info.getPopularity().split("Context: ");
|
||||
if (popParts.length > 1) {
|
||||
contextPrefix = popParts[1].split(" ")[0].toLowerCase();
|
||||
}
|
||||
}
|
||||
|
||||
LinkedHashSet<String> attempts = new LinkedHashSet<>();
|
||||
if (!contextPrefix.isEmpty()) {
|
||||
attempts.add(contextPrefix + " " + modelName);
|
||||
attempts.add(contextPrefix + " " + cleanName);
|
||||
}
|
||||
attempts.add(modelName);
|
||||
attempts.add(cleanName);
|
||||
attempts.add(baseName);
|
||||
|
||||
// Suche in offiziellen Repos (Top 100)
|
||||
for (String query : attempts) {
|
||||
if (query.length() < 3) continue;
|
||||
onStatusUpdate.accept(index, "🔍 Searching Official: " + query);
|
||||
if (fetchHuggingFaceUrl(info, index, query, true, onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
// Suche in Community Repos (Top 100)
|
||||
for (String query : attempts) {
|
||||
if (query.length() < 3) continue;
|
||||
onStatusUpdate.accept(index, "🔍 Searching Community: " + query);
|
||||
if (fetchHuggingFaceUrl(info, index, query, false, onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
// Civitai Fallback
|
||||
for (String query : attempts) {
|
||||
if (query.length() < 3) continue;
|
||||
onStatusUpdate.accept(index, "🔍 Searching Civitai: " + query);
|
||||
if (fetchCivitaiUrl(info, index, query, onStatusUpdate, onModelFound)) return;
|
||||
}
|
||||
|
||||
onStatusUpdate.accept(index, "❌ No trusted match found");
|
||||
}
|
||||
|
||||
private boolean validateAndSetUrl(ModelInfo info, int rowIndex, String url, String popPrefix,
|
||||
BiConsumer<Integer, String> onStatusUpdate,
|
||||
BiConsumer<Integer, ModelInfo> onModelFound) {
|
||||
if (url == null) return false;
|
||||
|
||||
// Normalize Hugging Face blob URLs to resolve URLs
|
||||
if (url.contains("huggingface.co") && url.contains("/blob/")) {
|
||||
url = url.replace("/blob/", "/resolve/");
|
||||
}
|
||||
|
||||
long size = getRemoteSize(url);
|
||||
if (size > 100) {
|
||||
info.setUrl(url);
|
||||
|
||||
@@ -26,7 +26,7 @@ public class ComfyModelAnalyzerTest {
|
||||
.orElse(null);
|
||||
|
||||
assertNotNull(flux);
|
||||
assertEquals("checkpoints", flux.getType());
|
||||
assertEquals("unet", flux.getType());
|
||||
assertEquals("https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors", flux.getUrl());
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ConfigService;
|
||||
import de.tki.comfymodels.service.impl.DefaultDownloadManager;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.io.File;
|
||||
import java.io.IOException;
|
||||
import java.io.OutputStream;
|
||||
import java.lang.reflect.Field;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CopyOnWriteArrayList;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
public class DownloadStatusTest {
|
||||
|
||||
private DefaultDownloadManager downloadManager;
|
||||
private HttpServer server;
|
||||
private String serverUrl;
|
||||
private Path tempDir;
|
||||
|
||||
@BeforeEach
|
||||
public void setup() throws Exception {
|
||||
tempDir = Files.createTempDirectory("downloader_test");
|
||||
downloadManager = new DefaultDownloadManager();
|
||||
|
||||
// Use real ConfigService
|
||||
ConfigService configService = new ConfigService();
|
||||
|
||||
Field field = DefaultDownloadManager.class.getDeclaredField("configService");
|
||||
field.setAccessible(true);
|
||||
field.set(downloadManager, configService);
|
||||
|
||||
// Setup Local Test Server
|
||||
server = HttpServer.create(new InetSocketAddress(0), 0);
|
||||
server.createContext("/test.safetensors", exchange -> {
|
||||
byte[] data = new byte[1024 * 1024]; // 1MB Test Data
|
||||
exchange.getResponseHeaders().add("Content-Length", String.valueOf(data.length));
|
||||
exchange.sendResponseHeaders(200, data.length);
|
||||
try (OutputStream os = exchange.getResponseBody()) {
|
||||
// Sende Daten langsam, damit wir Pause/Stop testen können
|
||||
for (int i = 0; i < 10; i++) {
|
||||
os.write(data, i * (data.length / 10), data.length / 10);
|
||||
os.flush();
|
||||
Thread.sleep(100);
|
||||
}
|
||||
} catch (InterruptedException ignored) {}
|
||||
});
|
||||
server.start();
|
||||
serverUrl = "http://localhost:" + server.getAddress().getPort() + "/test.safetensors";
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
public void tearDown() throws IOException {
|
||||
server.stop(0);
|
||||
// Clean temp files
|
||||
Files.walk(tempDir).map(Path::toFile).forEach(File::delete);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testFullStatusLifecycle() throws InterruptedException {
|
||||
List<String> capturedStatus = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch finishedLatch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "test_model.safetensors", serverUrl);
|
||||
List<ModelInfo> list = Collections.singletonList(info);
|
||||
boolean[] selected = {true};
|
||||
|
||||
downloadManager.startQueue(list, selected, tempDir.toString(),
|
||||
(idx, status) -> {
|
||||
capturedStatus.add(status);
|
||||
System.out.println("Status Update: " + status);
|
||||
},
|
||||
finishedLatch::countDown
|
||||
);
|
||||
|
||||
// 1. Prüfe auf Start/Download
|
||||
Thread.sleep(300);
|
||||
assertTrue(capturedStatus.stream().anyMatch(s -> s.contains("Downloading")), "Should show Downloading status");
|
||||
|
||||
// 2. Teste PAUSE
|
||||
System.out.println("--- Triggering PAUSE ---");
|
||||
downloadManager.togglePause();
|
||||
Thread.sleep(500);
|
||||
assertTrue(capturedStatus.contains("Paused"), "Should show Paused status");
|
||||
|
||||
// 3. Teste RESUME
|
||||
System.out.println("--- Triggering RESUME ---");
|
||||
downloadManager.togglePause();
|
||||
Thread.sleep(500);
|
||||
assertTrue(capturedStatus.contains("Resuming..."), "Should show Resuming status");
|
||||
|
||||
// 4. Teste STOP
|
||||
System.out.println("--- Triggering STOP ---");
|
||||
downloadManager.stop();
|
||||
|
||||
finishedLatch.await(5, TimeUnit.SECONDS);
|
||||
// Accept both "Stopped" or "Error: closed" (which is the result of closing the stream during stop)
|
||||
assertTrue(capturedStatus.stream().anyMatch(s -> s.contains("Stopped") || s.contains("Error: closed")),
|
||||
"Should show Stopped or closed status after stop(), but got: " + capturedStatus);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
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 static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
public class HiDreamReproductionTest {
|
||||
|
||||
@Test
|
||||
public void testHiDreamModelDetection() {
|
||||
String json = "{" +
|
||||
"\"id\": \"01d66ae9-78be-4a8d-b737-24eee5e1d447\"," +
|
||||
"\"nodes\": [" +
|
||||
" {\"id\": 55, \"type\": \"VAELoader\", \"widgets_values\": [\"ae.safetensors\"]}," +
|
||||
" {\"id\": 69, \"type\": \"UNETLoader\", \"widgets_values\": [\"hidream_i1_dev_bf16.safetensors\", \"default\"]}," +
|
||||
" {\"id\": 54, \"type\": \"QuadrupleCLIPLoader\", \"widgets_values\": [" +
|
||||
" \"clip_l_hidream.safetensors\", \"clip_g_hidream.safetensors\", " +
|
||||
" \"t5xxl_fp8_e4m3fn_scaled.safetensors\", \"llama_3.1_8b_instruct_fp8_scaled.safetensors\"" +
|
||||
" ]}," +
|
||||
" {\"id\": 72, \"type\": \"MarkdownNote\", \"widgets_values\": [\"## Official sampling settings\\n\\n### HiDream Full\\n* hidream_i1_full_fp16.safetensors\"]}" +
|
||||
"]" +
|
||||
"}";
|
||||
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
List<ModelInfo> models = analyzer.analyze(json, "workflow.json");
|
||||
|
||||
// Wir erwarten 7 Modelle: 1 UNET, 4 CLIPs, 1 VAE (aus Loadern) UND 1 UNET (aus der MarkdownNote)
|
||||
assertEquals(7, models.size(), "Should detect 7 models (including one from MarkdownNote)");
|
||||
|
||||
boolean foundUnet = false;
|
||||
int clipsFound = 0;
|
||||
boolean foundVae = false;
|
||||
|
||||
for (ModelInfo m : models) {
|
||||
if (m.getName().equals("hidream_i1_dev_bf16.safetensors")) {
|
||||
assertEquals("unet", m.getType());
|
||||
foundUnet = true;
|
||||
} else if (m.getName().equals("ae.safetensors")) {
|
||||
assertEquals("vae", m.getType());
|
||||
foundVae = true;
|
||||
} else if (m.getName().contains("clip") || m.getName().contains("t5xxl") || m.getName().contains("llama")) {
|
||||
assertEquals("clip", m.getType(), "Type for " + m.getName() + " should be clip");
|
||||
clipsFound++;
|
||||
}
|
||||
}
|
||||
|
||||
assertTrue(foundUnet, "UNET from Loader not found");
|
||||
assertTrue(foundVae, "VAE not found");
|
||||
assertEquals(4, clipsFound, "Should find 4 CLIP/TextEncoder models");
|
||||
|
||||
// Überprüfe ob der Kontext "hidream" erkannt wurde
|
||||
boolean contextDetected = false;
|
||||
for (ModelInfo m : models) {
|
||||
System.out.println("Model: " + m.getName() + " | Popularity: " + m.getPopularity());
|
||||
if (m.getPopularity() != null && m.getPopularity().toLowerCase().contains("hidream")) {
|
||||
contextDetected = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
assertTrue(contextDetected, "Context 'hidream' should be detected from workflow content");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
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 static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
public class QualityAssuranceTest {
|
||||
|
||||
@Test
|
||||
public void testFalsePositiveFiltering() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
|
||||
// Workflow mit "None" und "default" Werten in Widgets
|
||||
String json = "{" +
|
||||
"\"nodes\": [" +
|
||||
" {\"id\": 1, \"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"None\", \"default\"]}," +
|
||||
" {\"id\": 2, \"type\": \"LoraLoader\", \"widgets_values\": [\"undefined\", \"test.safetensors\"]}" +
|
||||
"]" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(json, "test.json");
|
||||
|
||||
// Es sollte nur "test.safetensors" gefunden werden
|
||||
assertEquals(1, models.size(), "Should only find one real model");
|
||||
assertEquals("test.safetensors", models.get(0).getName());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testMalformedJsonHandling() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
|
||||
// Kaputtes JSON
|
||||
String json = "{ \"nodes\": [ { \"id\": 1, \"type\": \"VAELoader\", \"widgets_values\": [\"ae.safetensors\""; // Fehlende Klammern
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(json, "broken.json");
|
||||
|
||||
// Sollte nicht abstürzen, sondern einfach 0 oder gefundene Modelle (je nach Parser-Robustheit) zurückgeben
|
||||
assertNotNull(models);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
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 ReproductionTest {
|
||||
|
||||
private final ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
|
||||
@Test
|
||||
public void testDeduplicationAndUrlFiltering() {
|
||||
String json = "{\n" +
|
||||
" \"nodes\": [\n" +
|
||||
" {\n" +
|
||||
" \"type\": \"MarkdownNote\",\n" +
|
||||
" \"widgets_values\": [\n" +
|
||||
" \"- [qwen_3_4b.safetensors](https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors)\"\n" +
|
||||
" ]\n" +
|
||||
" },\n" +
|
||||
" {\n" +
|
||||
" \"type\": \"CLIPLoader\",\n" +
|
||||
" \"properties\": {\n" +
|
||||
" \"models\": [\n" +
|
||||
" {\n" +
|
||||
" \"name\": \"qwen_3_4b.safetensors\",\n" +
|
||||
" \"url\": \"https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors\",\n" +
|
||||
" \"directory\": \"text_encoders\"\n" +
|
||||
" }\n" +
|
||||
" ]\n" +
|
||||
" }\n" +
|
||||
" }\n" +
|
||||
" ]\n" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(json, "test.json");
|
||||
|
||||
// Erwartung: Nur EIN Eintrag für qwen_3_4b.safetensors
|
||||
// Der Name darf NICHT die URL sein.
|
||||
long count = models.stream()
|
||||
.filter(m -> m.getName().toLowerCase().contains("qwen"))
|
||||
.count();
|
||||
|
||||
for (ModelInfo m : models) {
|
||||
System.out.println("Found: " + m.getName() + " [" + m.getType() + "] URL: " + m.getUrl());
|
||||
}
|
||||
|
||||
assertEquals(1, count, "Should only find one entry for qwen, not a second one with the URL as name");
|
||||
assertEquals("qwen_3_4b.safetensors", models.get(0).getName());
|
||||
assertEquals("text_encoders", models.get(0).getType());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ConfigService;
|
||||
import de.tki.comfymodels.service.impl.DefaultDownloadManager;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
public class DownloadSelectionTest {
|
||||
|
||||
@Test
|
||||
public void testOnlySelectedItemsAreProcessed() throws InterruptedException {
|
||||
DefaultDownloadManager downloadManager = new DefaultDownloadManager();
|
||||
|
||||
// Mocking ConfigService to avoid NPE
|
||||
ConfigService configService = new ConfigService();
|
||||
ReflectionTestUtils.setField(downloadManager, "configService", configService);
|
||||
|
||||
List<ModelInfo> models = new ArrayList<>();
|
||||
models.add(new ModelInfo("checkpoints", "model1.safetensors", "http://example.com/1"));
|
||||
models.add(new ModelInfo("checkpoints", "model2.safetensors", "http://example.com/2"));
|
||||
models.add(new ModelInfo("checkpoints", "model3.safetensors", "http://example.com/3"));
|
||||
|
||||
// Only first and third are selected
|
||||
boolean[] selected = {true, false, true};
|
||||
|
||||
ConcurrentHashMap<Integer, String> statusMap = new ConcurrentHashMap<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
downloadManager.startQueue(models, selected, "target/test_downloads", (idx, status) -> {
|
||||
statusMap.put(idx, status);
|
||||
}, () -> {
|
||||
latch.countDown();
|
||||
});
|
||||
|
||||
// Wait for completion
|
||||
assertTrue(latch.await(5, TimeUnit.SECONDS), "Download queue did not finish in time");
|
||||
|
||||
// Model 2 MUST be skipped
|
||||
assertEquals("Skipped (Not Selected)", statusMap.get(1), "Model 2 should have been skipped");
|
||||
|
||||
// Model 1 and 3 should have been attempted (likely Error because URL is fake, but NOT Skipped (Not Selected))
|
||||
assertTrue(statusMap.containsKey(0));
|
||||
assertTrue(!statusMap.get(0).equals("Skipped (Not Selected)"));
|
||||
assertTrue(statusMap.containsKey(2));
|
||||
assertTrue(!statusMap.get(2).equals("Skipped (Not Selected)"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testDynamicUncheckWhilePaused() throws InterruptedException {
|
||||
DefaultDownloadManager downloadManager = new DefaultDownloadManager();
|
||||
ConfigService configService = new ConfigService();
|
||||
ReflectionTestUtils.setField(downloadManager, "configService", configService);
|
||||
|
||||
List<ModelInfo> models = new ArrayList<>();
|
||||
// Use a dummy model
|
||||
models.add(new ModelInfo("checkpoints", "model1.safetensors", "http://invalid-url-for-test.com/file"));
|
||||
|
||||
boolean[] selected = {true};
|
||||
ConcurrentHashMap<Integer, String> statusMap = new ConcurrentHashMap<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
// Start in PAUSED state
|
||||
downloadManager.togglePause();
|
||||
|
||||
downloadManager.startQueue(models, selected, "target/test_downloads", (idx, status) -> {
|
||||
statusMap.put(idx, status);
|
||||
}, () -> {
|
||||
latch.countDown();
|
||||
});
|
||||
|
||||
// Give it a moment to reach the pause loop in startQueue or downloadWithResume
|
||||
Thread.sleep(300);
|
||||
|
||||
// At this point, it should be either "Idle" (not yet started) or "Paused"
|
||||
// But the important thing is: We uncheck it NOW while it's waiting
|
||||
downloadManager.updateSelection(new boolean[]{false});
|
||||
|
||||
// Resume the manager
|
||||
downloadManager.togglePause();
|
||||
|
||||
// Wait for completion (the thread should now see it's unchecked and finish)
|
||||
assertTrue(latch.await(5, TimeUnit.SECONDS), "Download queue did not finish in time");
|
||||
|
||||
// The status MUST indicate it was skipped or errored out (as long as it didn't finish)
|
||||
String finalStatus = statusMap.get(0);
|
||||
assertTrue(finalStatus != null && (finalStatus.contains("Skipped") || finalStatus.contains("ConnectException")),
|
||||
"Expected 'Skipped' or 'ConnectException' status but got: " + finalStatus);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.ComfyModelAnalyzer;
|
||||
import org.json.JSONObject;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
|
||||
public class ModelTypePathTest {
|
||||
|
||||
@Test
|
||||
public void testModelTypeMapping() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
|
||||
// Mock JSON representing different nodes
|
||||
String json = "{" +
|
||||
"\"1\": {\"class_type\": \"LoraLoader\", \"widgets_values\": [\"my_lora.safetensors\"]}," +
|
||||
"\"2\": {\"class_type\": \"ControlNetApply\", \"widgets_values\": [\"cn_model.pth\"]}," +
|
||||
"\"3\": {\"class_type\": \"CLIPLoader\", \"widgets_values\": [\"t5xxl.safetensors\"]}," +
|
||||
"\"4\": {\"class_type\": \"VAELoader\", \"widgets_values\": [\"vae_model.safetensors\"]}," +
|
||||
"\"5\": {\"class_type\": \"UNETLoader\", \"widgets_values\": [\"diffusion_model.safetensors\"]}," +
|
||||
"\"6\": {\"class_type\": \"GLIGENLoader\", \"widgets_values\": [\"gligen_box.safetensors\"]}," +
|
||||
"\"7\": {\"class_type\": \"CLIPVisionLoader\", \"widgets_values\": [\"clip_vision.safetensors\"]}," +
|
||||
"\"8\": {\"class_type\": \"AudioEncoderLoader\", \"widgets_values\": [\"audio.safetensors\"]}" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> results = analyzer.analyze(json, "test_workflow.json");
|
||||
results.forEach(r -> System.out.println("Found: " + r.getName() + " -> " + r.getType()));
|
||||
|
||||
assertPath(results, "my_lora.safetensors", "loras");
|
||||
assertPath(results, "cn_model.pth", "controlnet");
|
||||
assertPath(results, "t5xxl.safetensors", "clip");
|
||||
assertPath(results, "vae_model.safetensors", "vae");
|
||||
assertPath(results, "diffusion_model.safetensors", "unet");
|
||||
assertPath(results, "gligen_box.safetensors", "gligen");
|
||||
assertPath(results, "clip_vision.safetensors", "clip_vision");
|
||||
assertPath(results, "audio.safetensors", "audio_encoders");
|
||||
}
|
||||
|
||||
private void assertPath(List<ModelInfo> results, String name, String expectedType) {
|
||||
for (ModelInfo info : results) {
|
||||
if (info.getName().equals(name)) {
|
||||
assertEquals(expectedType, info.getType(), "Type mismatch for " + name);
|
||||
return;
|
||||
}
|
||||
}
|
||||
throw new AssertionError("Model " + name + " not found in results");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
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 QwenReproductionTest {
|
||||
|
||||
@Test
|
||||
public void testQwenMarkdownModelDetection() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
// Minimaler Workflow mit der MarkdownNote des Nutzers
|
||||
String workflow = "{\n" +
|
||||
" \"nodes\": [\n" +
|
||||
" {\n" +
|
||||
" \"id\": 99,\n" +
|
||||
" \"type\": \"MarkdownNote\",\n" +
|
||||
" \"widgets_values\": [\n" +
|
||||
" \"## Model links\\n\\n**LoRA**\\n\\n- [Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors](https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors)\\n\\n📂 loras/\\n├── 📂 loras/\\n│ └── Qwen-Image-Lightning-4steps-V1.0.safetensors\"\n" +
|
||||
" ]\n" +
|
||||
" }\n" +
|
||||
" ]\n" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(workflow, "test.json");
|
||||
|
||||
// Wir prüfen, ob beide Namen gefunden werden
|
||||
boolean foundLightning = models.stream().anyMatch(m -> m.getName().equals("Qwen-Image-Lightning-4steps-V1.0.safetensors"));
|
||||
boolean foundEdit = models.stream().anyMatch(m -> m.getName().equals("Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors"));
|
||||
|
||||
assertTrue(foundEdit, "Modell aus dem Link sollte gefunden werden");
|
||||
assertTrue(foundLightning, "Modell aus dem Text-Diagramm sollte gefunden werden");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testMarkdownLinkUrlExtraction() {
|
||||
ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer();
|
||||
String workflow = "{\n" +
|
||||
" \"nodes\": [\n" +
|
||||
" {\n" +
|
||||
" \"type\": \"MarkdownNote\",\n" +
|
||||
" \"widgets_values\": [\n" +
|
||||
" \"- [Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors](https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors)\"\n" +
|
||||
" ]\n" +
|
||||
" }\n" +
|
||||
" ]\n" +
|
||||
"}";
|
||||
|
||||
List<ModelInfo> models = analyzer.analyze(workflow, "test.json");
|
||||
|
||||
ModelInfo editModel = models.stream()
|
||||
.filter(m -> m.getName().equals("Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors"))
|
||||
.findFirst().orElse(null);
|
||||
|
||||
assertNotNull(editModel, "Modell sollte gefunden werden");
|
||||
assertEquals("https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors",
|
||||
editModel.getUrl(), "URL sollte aus dem Markdown-Link extrahiert werden");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
import de.tki.comfymodels.domain.ModelInfo;
|
||||
import de.tki.comfymodels.service.impl.*;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
|
||||
import java.io.File;
|
||||
import java.io.IOException;
|
||||
import java.io.OutputStream;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.Collections;
|
||||
import java.util.Comparator;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CopyOnWriteArrayList;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
public class StatusTransitionTest {
|
||||
|
||||
private DefaultDownloadManager downloadManager;
|
||||
private HttpServer server;
|
||||
private String serverUrl;
|
||||
private Path tempDir;
|
||||
private ConfigService configService;
|
||||
|
||||
@BeforeEach
|
||||
public void setup() throws Exception {
|
||||
tempDir = Files.createTempDirectory("transition_test");
|
||||
downloadManager = new DefaultDownloadManager();
|
||||
configService = new ConfigService();
|
||||
ReflectionTestUtils.setField(downloadManager, "configService", configService);
|
||||
|
||||
// Setup Local Test Server
|
||||
server = HttpServer.create(new InetSocketAddress(0), 0);
|
||||
server.createContext("/test.safetensors", exchange -> {
|
||||
if ("HEAD".equalsIgnoreCase(exchange.getRequestMethod())) {
|
||||
exchange.getResponseHeaders().add("Content-Length", "10000");
|
||||
exchange.sendResponseHeaders(200, -1);
|
||||
return;
|
||||
}
|
||||
byte[] data = new byte[10000];
|
||||
exchange.sendResponseHeaders(200, data.length);
|
||||
try (OutputStream os = exchange.getResponseBody()) {
|
||||
for (int i = 0; i < 5; i++) {
|
||||
os.write(data, i * 2000, 2000);
|
||||
os.flush();
|
||||
Thread.sleep(100);
|
||||
}
|
||||
} catch (Exception ignored) {}
|
||||
});
|
||||
|
||||
server.createContext("/auth-fail", exchange -> {
|
||||
exchange.sendResponseHeaders(401, -1);
|
||||
});
|
||||
|
||||
server.start();
|
||||
serverUrl = "http://localhost:" + server.getAddress().getPort();
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
public void tearDown() throws IOException {
|
||||
server.stop(0);
|
||||
if (Files.exists(tempDir)) {
|
||||
Files.walk(tempDir)
|
||||
.sorted(Comparator.reverseOrder())
|
||||
.map(Path::toFile)
|
||||
.forEach(File::delete);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Idle_To_Downloading_To_Finished() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
assertTrue(latch.await(5, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.stream().anyMatch(s -> s.contains("Downloading")), "Statuses: " + statuses);
|
||||
assertTrue(statuses.contains("✅ Finished"), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Idle_To_Searching_To_NotFound() throws InterruptedException {
|
||||
ModelSearchService searchService = new ModelSearchService();
|
||||
ConfigService cfg = new ConfigService();
|
||||
GeminiAIService gemini = new GeminiAIService();
|
||||
ModelListService list = new ModelListService();
|
||||
|
||||
ReflectionTestUtils.setField(searchService, "configService", cfg);
|
||||
ReflectionTestUtils.setField(searchService, "geminiService", gemini);
|
||||
ReflectionTestUtils.setField(searchService, "modelListService", list);
|
||||
ReflectionTestUtils.setField(gemini, "configService", cfg);
|
||||
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
ModelInfo info = new ModelInfo("checkpoints", "non_existent_model_xyz.safetensors", "MISSING");
|
||||
|
||||
searchService.searchOnline(Collections.singletonList(info), new boolean[]{true}, "context", "file.json",
|
||||
(idx, s) -> statuses.add(s),
|
||||
(idx, found) -> {},
|
||||
() -> {}
|
||||
);
|
||||
|
||||
// Wait a bit for search to progress (it will fail anyway as no API key is set)
|
||||
Thread.sleep(1500);
|
||||
assertTrue(statuses.stream().anyMatch(s -> s.contains("Searching") || s.contains("Scouting")), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Idle_To_SkippedNotSelected() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "skip_test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{false}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
assertTrue(latch.await(2, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.contains("Skipped (Not Selected)"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Downloading_To_Paused_To_Resuming() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "pause_test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
Thread.sleep(300);
|
||||
downloadManager.togglePause();
|
||||
Thread.sleep(500);
|
||||
assertTrue(statuses.contains("Paused"), "Statuses: " + statuses);
|
||||
|
||||
downloadManager.togglePause(); // Resume
|
||||
Thread.sleep(500);
|
||||
assertTrue(statuses.contains("Resuming..."), "Statuses: " + statuses);
|
||||
|
||||
downloadManager.stop();
|
||||
latch.await(2, TimeUnit.SECONDS);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Paused_To_SkippedUnchecked() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "uncheck_test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
Thread.sleep(300);
|
||||
downloadManager.togglePause();
|
||||
Thread.sleep(300);
|
||||
|
||||
downloadManager.updateSelection(new boolean[]{false});
|
||||
downloadManager.togglePause(); // Resume triggers check
|
||||
|
||||
assertTrue(latch.await(2, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.contains("Skipped (Unchecked)"), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Downloading_To_Stopped() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "stop_test.safetensors", serverUrl + "/test.safetensors");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
Thread.sleep(300);
|
||||
downloadManager.stop();
|
||||
|
||||
assertTrue(latch.await(2, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.contains("Stopped"), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Resuming_To_AuthRequired() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "auth.safetensors", serverUrl + "/auth-fail");
|
||||
|
||||
downloadManager.togglePause();
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
Thread.sleep(300);
|
||||
downloadManager.togglePause(); // Resume -> will hit 401
|
||||
|
||||
assertTrue(latch.await(2, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.contains("❌ Auth Required (Token?)"), "Statuses: " + statuses);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testTransition_Downloading_To_Error() throws InterruptedException {
|
||||
List<String> statuses = new CopyOnWriteArrayList<>();
|
||||
CountDownLatch latch = new CountDownLatch(1);
|
||||
|
||||
ModelInfo info = new ModelInfo("checkpoints", "error.safetensors", "http://localhost:1");
|
||||
downloadManager.startQueue(Collections.singletonList(info), new boolean[]{true}, tempDir.toString(),
|
||||
(idx, s) -> statuses.add(s), latch::countDown);
|
||||
|
||||
assertTrue(latch.await(5, TimeUnit.SECONDS));
|
||||
assertTrue(statuses.stream().anyMatch(s -> s.startsWith("Error:")), "Statuses: " + statuses);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package de.tki.comfymodels.service;
|
||||
|
||||
import de.tki.comfymodels.service.impl.ConfigService;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
import java.io.File;
|
||||
import java.nio.file.Path;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
public class VaultIntegrationTest {
|
||||
|
||||
@TempDir
|
||||
Path tempDir;
|
||||
|
||||
@Test
|
||||
public void testVaultSaveAndLoad() throws Exception {
|
||||
// Setup ConfigService with temp directory
|
||||
ConfigService configService = new ConfigService() {
|
||||
@Override
|
||||
public String getAppDataPath() {
|
||||
return tempDir.toString();
|
||||
}
|
||||
};
|
||||
|
||||
String password = "1234";
|
||||
String testKey = "4711";
|
||||
|
||||
// 1. Initial unlock (creates new vault)
|
||||
configService.unlock(password);
|
||||
|
||||
// 2. Set key and save
|
||||
configService.setGeminiApiKey(testKey);
|
||||
configService.save();
|
||||
|
||||
// 3. Create NEW instance to simulate restart
|
||||
ConfigService newConfigService = new ConfigService() {
|
||||
@Override
|
||||
public String getAppDataPath() {
|
||||
return tempDir.toString();
|
||||
}
|
||||
};
|
||||
|
||||
// 4. Unlock with correct password
|
||||
newConfigService.unlock(password);
|
||||
|
||||
// 5. Verify key
|
||||
assertEquals(testKey, newConfigService.getGeminiApiKey(), "Key should be 4711 after loading");
|
||||
|
||||
// 6. Verify wrong password fails
|
||||
ConfigService wrongPassService = new ConfigService() {
|
||||
@Override
|
||||
public String getAppDataPath() {
|
||||
return tempDir.toString();
|
||||
}
|
||||
};
|
||||
|
||||
assertThrows(Exception.class, () -> {
|
||||
wrongPassService.unlock("wrong_password");
|
||||
}, "Unlock should fail with wrong password");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user