diff --git a/.gitignore b/.gitignore index 0243213..abb1020 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ Thumbs.db # --- Project Specific --- # Sensitive API keys app_settings.json +settings.vault # User Model List (internal DB) uploaded_models.json diff --git a/README.md b/README.md index 73c2c43..c80d1d5 100644 --- a/README.md +++ b/README.md @@ -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.* diff --git a/pom.xml b/pom.xml index a79d352..a1d554f 100644 --- a/pom.xml +++ b/pom.xml @@ -5,12 +5,10 @@ 4.0.0 de.tki - comfymodeldownloader + comfyuimodel-downloader 1.0-SNAPSHOT - 11 - 11 UTF-8 @@ -42,6 +40,12 @@ 5.10.2 test + + org.springframework + spring-test + 6.1.5 + test + @@ -50,6 +54,9 @@ org.apache.maven.plugins maven-compiler-plugin 3.11.0 + + 11 + org.apache.maven.plugins diff --git a/run_downloader.bat b/run_downloader.bat new file mode 100644 index 0000000..b5f814b --- /dev/null +++ b/run_downloader.bat @@ -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 diff --git a/src/main/java/de/tki/comfymodels/Main.java b/src/main/java/de/tki/comfymodels/Main.java index 8387e68..3234001 100644 --- a/src/main/java/de/tki/comfymodels/Main.java +++ b/src/main/java/de/tki/comfymodels/Main.java @@ -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 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 files = (List) dtde.getTransferable().getTransferData(DataFlavor.javaFileListFlavor); + java.util.List files = (java.util.List) 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 = "" + + "

ComfyUIModel-Downloader Help

" + + "

Core Features

" + + "
    " + + "
  • Workflow Analysis: Drag and drop ComfyUI JSON or PNG files to extract required models.
    • " + + "
  • Deep Search: Uses AI and multi-platform searching (Hugging Face, Civitai) to find missing model URLs.
    • " + + "
  • Model List Integration: Match models against your own known-good local lists.
    • " + + "
  • Queue Management: Download multiple models simultaneously with pause/resume support.
    • " + + "
" + + "

API Keys & Security

" + + "
    " + + "
  • Gemini API Key: Required for the 'AI Scouting' feature. Gemini analyzes your workflow context " + + " to predict the best repositories for obscure models.
    • " + + "
  • Hugging Face Token: Recommended to avoid rate limits and to access gated repositories " + + " (like some FLUX or StabilityAI models).
    • " + + "
  • Local Storage (Vault): All keys and settings are stored locally on your machine " + + " in an encrypted file (settings.vault). They are never sent to any server except " + + " the official Google/Hugging Face APIs during requests.
    • " + + "
  • Password Protection: Your vault is protected by the password you choose at startup. " + + " Without this password, the API keys cannot be decrypted.
    • " + + "
" + + ""; + + 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 {} -} +} \ No newline at end of file diff --git a/src/main/java/de/tki/comfymodels/service/IDownloadManager.java b/src/main/java/de/tki/comfymodels/service/IDownloadManager.java index a65db5a..b6b2532 100644 --- a/src/main/java/de/tki/comfymodels/service/IDownloadManager.java +++ b/src/main/java/de/tki/comfymodels/service/IDownloadManager.java @@ -5,7 +5,8 @@ import java.util.List; import java.util.function.BiConsumer; public interface IDownloadManager { - void startQueue(List models, String baseDir, BiConsumer statusUpdater, Runnable onFinished); + void startQueue(List models, boolean[] selectedIndices, String baseDir, BiConsumer statusUpdater, Runnable onFinished); + void updateSelection(boolean[] selectedIndices); void togglePause(); void stop(); boolean isPaused(); diff --git a/src/main/java/de/tki/comfymodels/service/IModelSearchService.java b/src/main/java/de/tki/comfymodels/service/IModelSearchService.java index f068005..61e4a70 100644 --- a/src/main/java/de/tki/comfymodels/service/IModelSearchService.java +++ b/src/main/java/de/tki/comfymodels/service/IModelSearchService.java @@ -16,7 +16,8 @@ public interface IModelSearchService { */ void searchOnline(List modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName, BiConsumer onStatusUpdate, - BiConsumer onModelFound); + BiConsumer onModelFound, + Runnable onFinished); long getRemoteSize(String url); String formatSize(long bytes); diff --git a/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java b/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java index db9c2e5..78ecb85 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java @@ -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 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 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 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 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 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 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; } } diff --git a/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java b/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java index b0f2c7e..94fe5df 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java @@ -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; + } +} \ No newline at end of file diff --git a/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java b/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java index 56055c2..d55ce57 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java +++ b/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java @@ -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 completedIndices = java.util.concurrent.ConcurrentHashMap.newKeySet(); @Autowired private ConfigService configService; + private void safeUpdateStatus(int index, String status, BiConsumer 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 models, String baseDir, BiConsumer statusUpdater, Runnable onFinished) { + public void startQueue(List models, boolean[] selectedIndices, String baseDir, BiConsumer 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 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 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 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 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 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 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; } -} +} \ No newline at end of file diff --git a/src/main/java/de/tki/comfymodels/service/impl/EncryptionUtils.java b/src/main/java/de/tki/comfymodels/service/impl/EncryptionUtils.java new file mode 100644 index 0000000..9544fd3 --- /dev/null +++ b/src/main/java/de/tki/comfymodels/service/impl/EncryptionUtils.java @@ -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); + } +} \ No newline at end of file diff --git a/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java b/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java index ed9fa4e..8586d5f 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/GeminiAIService.java @@ -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" + diff --git a/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java b/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java index 6cddcb6..abb11bf 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java @@ -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 models = new ArrayList<>(); + @Autowired + private ConfigService configService; + @PostConstruct public void init() { loadFromStorage(); @@ -31,13 +35,14 @@ public class ModelListService { List 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(); } diff --git a/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java b/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java index c028f26..10a3787 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ModelSearchService.java @@ -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 modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName, - BiConsumer onStatusUpdate, - BiConsumer onModelFound) { - if (modelsToDownload == null) return; - - for (int i = 0; i < modelsToDownload.size(); i++) { - if (!selectedIndices[i]) continue; - - ModelInfo info = modelsToDownload.get(i); - final int index = i; - new Thread(() -> { - // Priority 1: User defined Model List - Optional manualMatch = modelListService.findByFilename(info.getName()); - if (manualMatch.isPresent() && !manualMatch.get().getUrl().equals("MISSING")) { - onStatusUpdate.accept(index, "📂 Found in Model List"); - if (validateAndSetUrl(info, index, manualMatch.get().getUrl(), "📂 USER DEFINED", onStatusUpdate, onModelFound)) return; - } - - onStatusUpdate.accept(index, "✨ Gemini Scouting..."); - String aiHint = geminiService.discoverBestRepo(info.getName(), fileName, workflowContext); - if (aiHint != null && !aiHint.equalsIgnoreCase("UNKNOWN")) { - if (aiHint.startsWith("http")) { - if (validateAndSetUrl(info, index, aiHint, "✨ AI DIRECT", onStatusUpdate, onModelFound)) return; - } - onStatusUpdate.accept(index, "🔍 Validating Repo: " + aiHint); - if (fetchHuggingFaceUrlInSpecificRepo(info, index, aiHint, onStatusUpdate, onModelFound)) return; - } - - String modelName = info.getName(); - String cleanName = modelName.replaceAll("(_fp8|_fp16|_bf16|_v\\d+|\\d+v|_fix|\\.safetensors|\\.sft|\\.ckpt)", ""); - String baseName = cleanName.split("[_\\- ]")[0]; - - String contextPrefix = fileName.toLowerCase().split("[_\\- ]")[0]; - if (contextPrefix.length() < 3) contextPrefix = ""; - - LinkedHashSet attempts = new LinkedHashSet<>(); - if (!contextPrefix.isEmpty()) { - attempts.add(contextPrefix + " " + modelName); - attempts.add(contextPrefix + " " + cleanName); - } - attempts.add(modelName); - attempts.add(cleanName); - attempts.add(baseName); - - // Suche in offiziellen Repos (Top 100) - for (String query : attempts) { - if (query.length() < 3) continue; - onStatusUpdate.accept(index, "🔍 Searching Official: " + query); - if (fetchHuggingFaceUrl(info, index, query, true, onStatusUpdate, onModelFound)) return; - } - - // Suche in Community Repos (Top 100) - for (String query : attempts) { - if (query.length() < 3) continue; - onStatusUpdate.accept(index, "🔍 Searching Community: " + query); - if (fetchHuggingFaceUrl(info, index, query, false, onStatusUpdate, onModelFound)) return; - } - - // Civitai Fallback - for (String query : attempts) { - if (query.length() < 3) continue; - onStatusUpdate.accept(index, "🔍 Searching Civitai: " + query); - if (fetchCivitaiUrl(info, index, query, onStatusUpdate, onModelFound)) return; - } - - onStatusUpdate.accept(index, "❌ No trusted match found"); - }).start(); + BiConsumer onStatusUpdate, + BiConsumer onModelFound, + Runnable onFinished) { + if (modelsToDownload == null) { + if (onFinished != null) onFinished.run(); + return; } + + List 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 onStatusUpdate, + BiConsumer onModelFound) { + // Priority 1: User defined Model List + Optional manualMatch = modelListService.findByFilename(info.getName()); + if (manualMatch.isPresent() && !manualMatch.get().getUrl().equals("MISSING")) { + onStatusUpdate.accept(index, "📂 Found in Model List"); + if (validateAndSetUrl(info, index, manualMatch.get().getUrl(), "📂 USER DEFINED", onStatusUpdate, onModelFound)) return; + } + + onStatusUpdate.accept(index, "✨ Gemini Scouting..."); + String aiHint = geminiService.discoverBestRepo(info.getName(), fileName, workflowContext); + if (aiHint != null && !aiHint.equalsIgnoreCase("UNKNOWN")) { + if (aiHint.startsWith("http")) { + if (validateAndSetUrl(info, index, aiHint, "✨ AI DIRECT", onStatusUpdate, onModelFound)) return; + } + onStatusUpdate.accept(index, "🔍 Validating Repo: " + aiHint); + if (fetchHuggingFaceUrlInSpecificRepo(info, index, aiHint, onStatusUpdate, onModelFound)) return; + } + + String modelName = info.getName(); + String cleanName = modelName.replaceAll("(_fp8|_fp16|_bf16|_v\\d+|\\d+v|_fix|\\.safetensors|\\.sft|\\.ckpt)", ""); + String[] 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 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 onStatusUpdate, BiConsumer 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); diff --git a/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java b/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java index d269745..4750ea4 100644 --- a/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java +++ b/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java @@ -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()); } diff --git a/src/test/java/de/tki/comfymodels/DownloadStatusTest.java b/src/test/java/de/tki/comfymodels/DownloadStatusTest.java new file mode 100644 index 0000000..4e53638 --- /dev/null +++ b/src/test/java/de/tki/comfymodels/DownloadStatusTest.java @@ -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 capturedStatus = new CopyOnWriteArrayList<>(); + CountDownLatch finishedLatch = new CountDownLatch(1); + + ModelInfo info = new ModelInfo("checkpoints", "test_model.safetensors", serverUrl); + List 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); + } +} diff --git a/src/test/java/de/tki/comfymodels/HiDreamReproductionTest.java b/src/test/java/de/tki/comfymodels/HiDreamReproductionTest.java new file mode 100644 index 0000000..8594d0e --- /dev/null +++ b/src/test/java/de/tki/comfymodels/HiDreamReproductionTest.java @@ -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 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"); + } +} diff --git a/src/test/java/de/tki/comfymodels/QualityAssuranceTest.java b/src/test/java/de/tki/comfymodels/QualityAssuranceTest.java new file mode 100644 index 0000000..9fbad1f --- /dev/null +++ b/src/test/java/de/tki/comfymodels/QualityAssuranceTest.java @@ -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 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 models = analyzer.analyze(json, "broken.json"); + + // Sollte nicht abstürzen, sondern einfach 0 oder gefundene Modelle (je nach Parser-Robustheit) zurückgeben + assertNotNull(models); + } +} diff --git a/src/test/java/de/tki/comfymodels/ReproductionTest.java b/src/test/java/de/tki/comfymodels/ReproductionTest.java new file mode 100644 index 0000000..d28890c --- /dev/null +++ b/src/test/java/de/tki/comfymodels/ReproductionTest.java @@ -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 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()); + } +} diff --git a/src/test/java/de/tki/comfymodels/service/DownloadSelectionTest.java b/src/test/java/de/tki/comfymodels/service/DownloadSelectionTest.java new file mode 100644 index 0000000..b28d6c8 --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/DownloadSelectionTest.java @@ -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 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 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 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 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); + } +} diff --git a/src/test/java/de/tki/comfymodels/service/ModelTypePathTest.java b/src/test/java/de/tki/comfymodels/service/ModelTypePathTest.java new file mode 100644 index 0000000..7e87de0 --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/ModelTypePathTest.java @@ -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 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 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"); + } +} diff --git a/src/test/java/de/tki/comfymodels/service/QwenReproductionTest.java b/src/test/java/de/tki/comfymodels/service/QwenReproductionTest.java new file mode 100644 index 0000000..07dbf6d --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/QwenReproductionTest.java @@ -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 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 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"); + } +} diff --git a/src/test/java/de/tki/comfymodels/service/StatusTransitionTest.java b/src/test/java/de/tki/comfymodels/service/StatusTransitionTest.java new file mode 100644 index 0000000..13c7b58 --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/StatusTransitionTest.java @@ -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 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 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 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 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 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 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 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 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); + } +} diff --git a/src/test/java/de/tki/comfymodels/service/VaultIntegrationTest.java b/src/test/java/de/tki/comfymodels/service/VaultIntegrationTest.java new file mode 100644 index 0000000..7f17422 --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/VaultIntegrationTest.java @@ -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"); + } +} \ No newline at end of file