From 826f0aa2b53c6301be4c6a35294610b754757f1f Mon Sep 17 00:00:00 2001 From: Thomas Date: Thu, 30 Apr 2026 21:11:24 +0200 Subject: [PATCH] feat: restore missing Import Model List button and add REST bridge for ComfyUI integration --- README.md | 9 + src/main/java/de/tki/comfymodels/Main.java | 833 ++++++++++-------- .../service/impl/ConfigService.java | 13 + .../service/impl/DefaultDownloadManager.java | 73 +- .../service/impl/RestBridgeService.java | 125 +++ .../service/InstallationPathTest.java | 88 ++ .../service/RestBridgeServiceTest.java | 53 ++ 7 files changed, 775 insertions(+), 419 deletions(-) create mode 100644 src/main/java/de/tki/comfymodels/service/impl/RestBridgeService.java create mode 100644 src/test/java/de/tki/comfymodels/service/InstallationPathTest.java create mode 100644 src/test/java/de/tki/comfymodels/service/RestBridgeServiceTest.java diff --git a/README.md b/README.md index 62704cb..3ac339a 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,15 @@ Stop wasting time manually moving files! This tool automates your model manageme * **Background Mode:** Keep the downloader running in the system tray. * **Vault Security:** Encrypted storage for your API keys (Gemini, Hugging Face). +## 🔌 ComfyUI Integration + +To send workflows directly from ComfyUI to the downloader with a single click: + +1. Copy or symlink the folder `comfyui-model-downloader` into your `ComfyUI/custom_nodes/` directory. +2. Restart ComfyUI. +3. You will now see a **🚀 Downloader** button in the ComfyUI menu (or under the download icon in the new UI v2). +4. Make sure the Java application is running to receive the workflow. + ## 🚀 Getting Started ### Prerequisites diff --git a/src/main/java/de/tki/comfymodels/Main.java b/src/main/java/de/tki/comfymodels/Main.java index ac0556b..ebef547 100644 --- a/src/main/java/de/tki/comfymodels/Main.java +++ b/src/main/java/de/tki/comfymodels/Main.java @@ -9,6 +9,7 @@ import de.tki.comfymodels.service.IModelSearchService; import de.tki.comfymodels.service.impl.ConfigService; import de.tki.comfymodels.service.impl.GeminiAIService; import de.tki.comfymodels.service.impl.ModelListService; +import de.tki.comfymodels.service.impl.ModelHashRegistry; import com.formdev.flatlaf.FlatDarkLaf; import com.formdev.flatlaf.FlatLightLaf; import com.formdev.flatlaf.FlatLaf; @@ -37,6 +38,7 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; +import java.nio.file.StandardCopyOption; import java.util.List; import java.util.ArrayList; import java.util.Map; @@ -50,6 +52,7 @@ public class Main extends JFrame { private final IWorkflowService workflowService; private final IModelSearchService searchService; private final IModelValidator modelValidator; + private final de.tki.comfymodels.service.impl.RestBridgeService restBridge; @Autowired private ConfigService configService; @@ -61,9 +64,10 @@ public class Main extends JFrame { private ModelListService modelListService; @Autowired - private de.tki.comfymodels.service.impl.ModelHashRegistry hashRegistry; + private ModelHashRegistry hashRegistry; private JTextField modelsPathField; + private JTextField comfyPathField; private JPasswordField geminiKeyField; private JPasswordField hfTokenField; private JCheckBox backgroundCheck; @@ -79,32 +83,49 @@ public class Main extends JFrame { public Main(IModelAnalyzer analyzer, IDownloadManager downloadManager, IWorkflowService workflowService, IModelSearchService searchService, - IModelValidator modelValidator) { + IModelValidator modelValidator, de.tki.comfymodels.service.impl.RestBridgeService restBridge) { this.analyzer = analyzer; this.downloadManager = downloadManager; this.workflowService = workflowService; this.searchService = searchService; this.modelValidator = modelValidator; + this.restBridge = restBridge; } public void launch(String[] args) { - boolean backgroundArg = false; - for (String arg : args) { - if ("--background".equalsIgnoreCase(arg)) { - backgroundArg = true; - break; - } + // Initialize REST Bridge consumer EARLY + restBridge.setWorkflowConsumer(workflowJson -> { + SwingUtilities.invokeLater(() -> { + if (jsonInputArea != null) { + jsonInputArea.setText(workflowJson); + currentFileName = "remote_workflow.json"; + analyzeJsonContent(); + } + setVisible(true); + toFront(); + requestFocus(); + }); + }); + + // Load settings to get/generate the API Token + if (configService.isUnlocked()) { + restBridge.setApiToken(configService.getApiToken()); } + restBridge.startServer(); if (!promptForPassword()) { System.exit(0); } // Apply theme before UI initialization - if (configService.isDarkMode()) { - FlatDarkLaf.setup(); - } else { - FlatLightLaf.setup(); + try { + if (configService.isDarkMode()) { + FlatDarkLaf.setup(); + } else { + FlatLightLaf.setup(); + } + } catch (Exception e) { + System.err.println("Theme setup failed: " + e.getMessage()); } FlatLaf.updateUI(); @@ -116,25 +137,30 @@ public class Main extends JFrame { UIManager.put("ProgressBar.arc", 8); SwingUtilities.invokeLater(() -> { - initUI(); - setupTrayIcon(); - loadSettingsIntoUI(); - updateAiModelDisplay(); - - // Add WindowListener to handle background mode - addWindowListener(new java.awt.event.WindowAdapter() { - @Override - public void windowClosing(java.awt.event.WindowEvent e) { - if (configService.isBackgroundModeEnabled()) { - setVisible(false); - } else { - downloadManager.stop(); - System.exit(0); + try { + initUI(); + setupTrayIcon(); + loadSettingsIntoUI(); + updateAiModelDisplay(); + + // Add WindowListener to handle background mode + addWindowListener(new java.awt.event.WindowAdapter() { + @Override + public void windowClosing(java.awt.event.WindowEvent e) { + if (configService.isBackgroundModeEnabled()) { + setVisible(false); + } else { + downloadManager.stop(); + System.exit(0); + } } - } - }); - - setVisible(true); + }); + + setVisible(true); + } catch (Exception e) { + e.printStackTrace(); + JOptionPane.showMessageDialog(null, "Critical UI Error: " + e.getMessage()); + } }); } @@ -142,7 +168,6 @@ public class Main extends JFrame { if (!SystemTray.isSupported()) return; SystemTray tray = SystemTray.getSystemTray(); - // Use a simple colored square as fallback (Pink for ComfyUI style) BufferedImage image = new BufferedImage(16, 16, BufferedImage.TYPE_INT_RGB); Graphics2D g2 = image.createGraphics(); g2.setColor(new Color(255, 105, 180)); // Hot Pink @@ -173,29 +198,6 @@ public class Main extends JFrame { } } - private void updatePendingDownloadsPersistence() { - if (modelsToDownload == null) return; - org.json.JSONArray arr = new org.json.JSONArray(); - for (int i = 0; i < modelsToDownload.size(); i++) { - // Only save if not finished - String status = ""; - if (tableModel != null && tableModel.getRowCount() > i) { - status = (String) tableModel.getValueAt(i, 7); - } - - if (!status.contains("✅ Finished") && !status.contains("Already exists")) { - ModelInfo info = modelsToDownload.get(i); - org.json.JSONObject obj = new org.json.JSONObject(); - obj.put("name", info.getName()); - obj.put("url", info.getUrl()); - obj.put("type", info.getType()); - obj.put("save_path", info.getSave_path()); - arr.put(obj); - } - } - configService.savePendingDownloads(arr.toString()); - } - private boolean promptForPassword() { while (true) { JPanel panel = new JPanel(new BorderLayout(5, 5)); @@ -203,7 +205,6 @@ public class Main extends JFrame { 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); @@ -217,6 +218,7 @@ public class Main extends JFrame { try { configService.unlock(pass); + restBridge.setApiToken(configService.getApiToken()); return true; } catch (Exception e) { JOptionPane.showMessageDialog(null, "Unlock Failed: " + e.getMessage(), "Error", JOptionPane.ERROR_MESSAGE); @@ -225,30 +227,13 @@ public class Main extends JFrame { } 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) + ")"); - } - if (backgroundCheck != null) { - backgroundCheck.setSelected(configService.isBackgroundModeEnabled()); - } - if (shutdownCheck != null) { - shutdownCheck.setSelected(configService.isShutdownAfterDownloadEnabled()); - } - if (darkCheck != null) { - darkCheck.setSelected(configService.isDarkMode()); - } + if (modelsPathField != null) modelsPathField.setText(configService.getModelsPath()); + if (comfyPathField != null) comfyPathField.setText(configService.getComfyUIPath()); + if (geminiKeyField != null) geminiKeyField.setText(configService.getGeminiApiKey()); + if (hfTokenField != null) hfTokenField.setText(configService.getHfToken()); + if (backgroundCheck != null) backgroundCheck.setSelected(configService.isBackgroundModeEnabled()); + if (shutdownCheck != null) shutdownCheck.setSelected(configService.isShutdownAfterDownloadEnabled()); + if (darkCheck != null) darkCheck.setSelected(configService.isDarkMode()); } private void updateAiModelDisplay() { @@ -263,15 +248,15 @@ public class Main extends JFrame { setSize(1450, 1000); setDefaultCloseOperation(JFrame.DO_NOTHING_ON_CLOSE); - // Enable FlatLaf window decorations for this frame getRootPane().putClientProperty("flatlaf.useWindowDecorations", true); JPanel mainContainer = new JPanel(new BorderLayout(10, 10)); mainContainer.setBorder(BorderFactory.createEmptyBorder(10, 10, 10, 10)); setContentPane(mainContainer); - JPanel settingsPanel = new JPanel(new GridLayout(4, 1, 5, 5)); + JPanel settingsPanel = new JPanel(new GridLayout(5, 1, 5, 5)); + // Row 1: Models Directory JPanel pathRow = new JPanel(new BorderLayout()); pathRow.setBorder(BorderFactory.createTitledBorder("Models Directory")); modelsPathField = new JTextField(); @@ -279,31 +264,30 @@ public class Main extends JFrame { 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) { - String selectedPath = chooser.getSelectedFile().getAbsolutePath(); - modelsPathField.setText(selectedPath); - configService.setModelsPath(selectedPath); + modelsPathField.setText(chooser.getSelectedFile().getAbsolutePath()); + configService.setModelsPath(modelsPathField.getText()); } }); - JButton savePathBtn = new JButton("Save Path"); savePathBtn.addActionListener(e -> { configService.setModelsPath(modelsPathField.getText()); - statusLabel.setText("Models path saved to vault."); + statusLabel.setText("Models path saved."); }); - JPanel pathButtons = new JPanel(new FlowLayout(FlowLayout.RIGHT, 0, 0)); pathButtons.add(savePathBtn); pathButtons.add(browsePathBtn); - pathRow.add(modelsPathField, BorderLayout.CENTER); pathRow.add(pathButtons, BorderLayout.EAST); + // Row 2: ComfyUI Installation (Simplified in Main) + JPanel comfyRow = new JPanel(new BorderLayout()); + comfyRow.setBorder(BorderFactory.createTitledBorder("ComfyUI Integration")); + JButton openInstallDialogBtn = new JButton("🚀 Install/Update Bridge..."); + openInstallDialogBtn.addActionListener(e -> showInstallationDialog()); + comfyRow.add(openInstallDialogBtn, BorderLayout.CENTER); + + // Row 3: Gemini API Key JPanel apiRow = new JPanel(new BorderLayout()); apiRow.setBorder(BorderFactory.createTitledBorder("Gemini AI API Key (Optional)")); geminiKeyField = new JPasswordField(); @@ -315,6 +299,7 @@ public class Main extends JFrame { apiRow.add(geminiKeyField, BorderLayout.CENTER); apiRow.add(saveGeminiBtn, BorderLayout.EAST); + // Row 4: Hugging Face Token JPanel hfRow = new JPanel(new BorderLayout()); hfRow.setBorder(BorderFactory.createTitledBorder("Hugging Face Access Token (Optional)")); hfTokenField = new JPasswordField(); @@ -327,10 +312,11 @@ public class Main extends JFrame { hfRow.add(hfTokenField, BorderLayout.CENTER); hfRow.add(saveHfBtn, BorderLayout.EAST); + // Row 5: Options & Info JPanel infoRow = new JPanel(new FlowLayout(FlowLayout.LEFT)); - JButton helpBtn = new JButton("ℹ Help & Security Information"); + activeAiModelLabel = new JLabel("Active AI: Loading..."); + JButton helpBtn = new JButton("ℹ Help"); helpBtn.addActionListener(e -> showHelpDialog()); - JButton verifyBtn = new JButton("🔍 Fast Verify (Corruption)"); verifyBtn.addActionListener(e -> verifyLocalModels(false)); @@ -345,28 +331,19 @@ public class Main extends JFrame { } }); - backgroundCheck = new JCheckBox("Stay in Background on Close"); + backgroundCheck = new JCheckBox("Stay in Background"); backgroundCheck.addActionListener(e -> configService.setBackgroundModeEnabled(backgroundCheck.isSelected())); - shutdownCheck = new JCheckBox("Shutdown after Queue"); shutdownCheck.addActionListener(e -> configService.setShutdownAfterDownloadEnabled(shutdownCheck.isSelected())); - - // Most important feature (irony): Dark mode for Reddit... or at least for this tool. darkCheck = new JCheckBox("Dark Mode"); darkCheck.addActionListener(e -> { - boolean isDark = darkCheck.isSelected(); - configService.setDarkMode(isDark); - + configService.setDarkMode(darkCheck.isSelected()); FlatAnimatedLafChange.showSnapshot(); - if (isDark) { - FlatDarkLaf.setup(); - } else { - FlatLightLaf.setup(); - } + if (darkCheck.isSelected()) FlatDarkLaf.setup(); else FlatLightLaf.setup(); FlatLaf.updateUI(); FlatAnimatedLafChange.hideSnapshotWithAnimation(); }); - + infoRow.add(new JSeparator(JSeparator.VERTICAL)); infoRow.add(helpBtn); infoRow.add(verifyBtn); infoRow.add(optimizeBtn); @@ -375,6 +352,7 @@ public class Main extends JFrame { infoRow.add(darkCheck); settingsPanel.add(pathRow); + settingsPanel.add(comfyRow); settingsPanel.add(apiRow); settingsPanel.add(hfRow); settingsPanel.add(infoRow); @@ -385,7 +363,6 @@ 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 -> { @@ -394,27 +371,23 @@ public class Main extends JFrame { if (fd.getFile() != null) loadFile(new File(fd.getDirectory(), fd.getFile())); }); - JButton importModelListBtn = new JButton("Import Model List (JSON)..."); + JButton importModelListBtn = new JButton("Import Model List..."); importModelListBtn.addActionListener(e -> { FileDialog fd = new FileDialog(this, "Select Model List JSON", FileDialog.LOAD); fd.setVisible(true); if (fd.getFile() != null) { try { modelListService.importJson(new File(fd.getDirectory(), fd.getFile())); - JOptionPane.showMessageDialog(this, "Model list imported successfully!"); + JOptionPane.showMessageDialog(this, "Model list imported successfully! (" + modelListService.getModels().size() + " models)"); analyzeJsonContent(); } catch (Exception ex) { - JOptionPane.showMessageDialog(this, "Error: " + ex.getMessage()); + JOptionPane.showMessageDialog(this, "Error: " + ex.getMessage()); } } }); - JButton analyzeBtn = new JButton("Deep Search (Strict Validation)"); - analyzeBtn.addActionListener(e -> { - analyzeJsonContent(); - searchMissingOnline(); - }); - + JButton analyzeBtn = new JButton("Deep Search"); + analyzeBtn.addActionListener(e -> { analyzeJsonContent(); searchMissingOnline(); }); jsonButtons.add(loadJsonBtn); jsonButtons.add(importModelListBtn); jsonButtons.add(analyzeBtn); @@ -427,26 +400,10 @@ public class Main extends JFrame { @Override public boolean isCellEditable(int r, int c) { return c == 0; } }; tableModel.addTableModelListener(e -> { - if (e.getType() == TableModelEvent.UPDATE && e.getColumn() == 0) { - updateDownloadManagerSelection(); - } + if (e.getType() == TableModelEvent.UPDATE && e.getColumn() == 0) updateDownloadManagerSelection(); }); JTable modelTable = new JTable(tableModel); modelTable.putClientProperty("terminateEditOnFocusLost", Boolean.TRUE); - - // Immediate termination of the editor on click to apply the value - 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(150); - modelTable.getColumnModel().getColumn(5).setMinWidth(150); JScrollPane tableScroll = new JScrollPane(modelTable); JPanel tablePanel = new JPanel(new BorderLayout()); tablePanel.setBorder(BorderFactory.createTitledBorder("Identified Models")); @@ -458,64 +415,28 @@ public class Main extends JFrame { JPanel bottomPanel = new JPanel(new BorderLayout()); statusLabel = new JLabel("Ready"); - + + JPanel progressPanel = new JPanel(new GridLayout(2, 1)); + activeAiModelLabel.setFont(new Font("SansSerif", Font.ITALIC, 11)); + activeAiModelLabel.setForeground(Color.GRAY); + progressPanel.add(statusLabel); + progressPanel.add(activeAiModelLabel); + 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); - updatePendingDownloadsPersistence(); - downloadManager.startQueue(modelsToDownload, selected, modelsPathField.getText(), - (idx, status) -> SwingUtilities.invokeLater(() -> { - tableModel.setValueAt(status, idx, 7); - if (status.contains("✅ Finished") || status.contains("Already exists")) { - updatePendingDownloadsPersistence(); - } - }), - () -> SwingUtilities.invokeLater(() -> { - downloadButton.setEnabled(true); - statusLabel.setText("Queue finished."); - updatePendingDownloadsPersistence(); - - if (configService.isShutdownAfterDownloadEnabled()) { - performSystemShutdown(); - } - }) - ); - }); - + downloadButton.addActionListener(e -> startDownloadQueue()); 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); + bottomPanel.add(progressPanel, BorderLayout.CENTER); bottomPanel.add(actionButtons, BorderLayout.EAST); @@ -524,125 +445,311 @@ public class Main extends JFrame { mainContainer.add(bottomPanel, BorderLayout.SOUTH); } - private void performSystemShutdown() { - String os = System.getProperty("os.name").toLowerCase(); - try { - if (os.contains("win")) { - Runtime.getRuntime().exec("shutdown /s /t 60"); - JOptionPane.showMessageDialog(this, "The system will shut down in 60 seconds. You can cancel this with 'shutdown /a' in the terminal.", "System Shutdown Scheduled", JOptionPane.WARNING_MESSAGE); - } else if (os.contains("nix") || os.contains("nux") || os.contains("mac")) { - Runtime.getRuntime().exec("shutdown -h +1"); - JOptionPane.showMessageDialog(this, "The system will shut down in 1 minute. You can cancel this with 'shutdown -c'.", "System Shutdown Scheduled", JOptionPane.WARNING_MESSAGE); - } - } catch (IOException e) { - JOptionPane.showMessageDialog(this, "Could not trigger system shutdown: " + e.getMessage(), "Error", JOptionPane.ERROR_MESSAGE); - } - } - - private void updateDownloadManagerSelection() { + private void startDownloadQueue() { int rowCount = tableModel.getRowCount(); + if (rowCount == 0) return; boolean[] selected = new boolean[rowCount]; for (int i = 0; i < rowCount; i++) { selected[i] = (Boolean) tableModel.getValueAt(i, 0); + String url = (String) tableModel.getValueAt(i, 6); + if (url != null && !url.equals("MISSING")) modelsToDownload.get(i).setUrl(url); } - downloadManager.updateSelection(selected); + downloadButton.setEnabled(false); + downloadManager.startQueue(modelsToDownload, selected, modelsPathField.getText(), + (idx, status) -> SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, idx, 7)), + () -> SwingUtilities.invokeLater(() -> { + downloadButton.setEnabled(true); + statusLabel.setText("Queue finished."); + if (configService.isShutdownAfterDownloadEnabled()) performSystemShutdown(); + }) + ); } - private void setupDragAndDrop(JTextArea area) { - new DropTarget(area, new DropTargetListener() { - public void dragEnter(DropTargetDragEvent dtde) {} - public void dragOver(DropTargetDragEvent dtde) {} - public void dropActionChanged(DropTargetDragEvent dtde) {} - public void dragExit(DropTargetEvent dte) {} - public void drop(DropTargetDropEvent dtde) { - try { - dtde.acceptDrop(DnDConstants.ACTION_COPY); - java.util.List files = (java.util.List) dtde.getTransferable().getTransferData(DataFlavor.javaFileListFlavor); - if (!files.isEmpty()) loadFile(files.get(0)); - } catch (Exception e) {} + private void showInstallationDialog() { + JDialog dialog = new JDialog(this, "ComfyUI Bridge Installation", true); + dialog.setLayout(new BorderLayout()); + dialog.setSize(650, 360); + dialog.setLocationRelativeTo(this); + + JPanel content = new JPanel(new GridBagLayout()); + content.setBorder(BorderFactory.createEmptyBorder(20, 20, 20, 20)); + GridBagConstraints gbc = new GridBagConstraints(); + gbc.fill = GridBagConstraints.HORIZONTAL; + gbc.gridx = 0; gbc.gridy = 0; gbc.weightx = 1.0; + gbc.insets = new Insets(0, 0, 15, 0); + + // Explanation text + JTextArea infoArea = new JTextArea( + "This installer will set up the ComfyUI-Model-Downloader bridge.\n\n" + + "1. Select your ComfyUI main directory.\n" + + "2. Old or conflicting bridge files will be cleaned up.\n" + + "3. A link or copy of the UI extension will be created.\n\n" + + "Important: Restart ComfyUI after the installation is finished." + ); + infoArea.setEditable(false); + infoArea.setFocusable(false); + infoArea.setBackground(content.getBackground()); + infoArea.setLineWrap(true); + infoArea.setWrapStyleWord(true); + content.add(infoArea, gbc); + + // Path selection header + gbc.gridy++; + gbc.insets = new Insets(0, 0, 5, 0); + content.add(new JLabel("ComfyUI Main Directory:"), gbc); + + // Path selection row + gbc.gridy++; + gbc.insets = new Insets(0, 0, 10, 0); + JPanel pathRow = new JPanel(new BorderLayout(5, 0)); + JTextField pathField = new JTextField(configService.getComfyUIPath()); + JButton browseBtn = new JButton("Browse..."); + browseBtn.addActionListener(e -> { + JFileChooser chooser = new JFileChooser(); + chooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY); + if (chooser.showOpenDialog(dialog) == JFileChooser.APPROVE_OPTION) { + pathField.setText(chooser.getSelectedFile().getAbsolutePath()); } }); + pathRow.add(pathField, BorderLayout.CENTER); + pathRow.add(browseBtn, BorderLayout.EAST); + content.add(pathRow, gbc); + + // Spacer to push everything to the top + gbc.gridy++; + gbc.weighty = 1.0; + content.add(new JPanel(), gbc); + + // Action buttons + JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT)); + JButton cancelBtn = new JButton("Cancel"); + cancelBtn.addActionListener(e -> dialog.dispose()); + JButton installBtn = new JButton("🚀 Start Installation"); + installBtn.addActionListener(e -> { + String selectedPath = pathField.getText().trim(); + if (selectedPath.isEmpty()) { + JOptionPane.showMessageDialog(dialog, "Please select a path first."); + return; + } + installComfyUIBridge(selectedPath, dialog); + }); + buttonPanel.add(cancelBtn); + buttonPanel.add(installBtn); + + dialog.add(content, BorderLayout.CENTER); + dialog.add(buttonPanel, BorderLayout.SOUTH); + dialog.setVisible(true); + } + + private void installComfyUIBridge(String comfyPath, JDialog parentDialog) { + File inputPath = new File(comfyPath); + if (!inputPath.exists()) { + JOptionPane.showMessageDialog(parentDialog, "The provided path does not exist: " + comfyPath, "Error", JOptionPane.ERROR_MESSAGE); + return; + } + + // Systematically search for custom_nodes + File customNodesDir = null; + + // Candidate 0: Specific Desktop App nested structure (High Priority) + File c0 = new File(inputPath, "resources" + File.separator + "ComfyUI" + File.separator + "custom_nodes"); + // Candidate 1: inputPath/custom_nodes (Standard) + File c1 = new File(inputPath, "custom_nodes"); + // Candidate 2: inputPath/ComfyUI/custom_nodes (Portable Root) + File c2 = new File(inputPath, "ComfyUI" + File.separator + "custom_nodes"); + // Candidate 3: inputPath itself (User selected custom_nodes directly) + File c3 = inputPath; + + if (c0.exists() && c0.isDirectory()) customNodesDir = c0; + else if (c1.exists() && c1.isDirectory()) customNodesDir = c1; + else if (c2.exists() && c2.isDirectory()) customNodesDir = c2; + else if (c3.getName().equalsIgnoreCase("custom_nodes") && c3.isDirectory()) customNodesDir = c3; + + // --- IMPROVEMENT: DEEP SEARCH (if candidates fail) --- + if (customNodesDir == null) { + System.out.println("Standard candidates failed. Starting deep search in: " + inputPath.getAbsolutePath()); + customNodesDir = findCustomNodesDeep(inputPath, 0); + } + + if (customNodesDir == null) { + String msg = "Could not find 'custom_nodes' folder in the selected directory.\n\n" + + "Checked locations (among others):\n" + + "- " + c0.getAbsolutePath() + "\n" + + "- " + c1.getAbsolutePath() + "\n" + + "- " + c2.getAbsolutePath() + "\n\n" + + "Please ensure you select the folder that contains 'custom_nodes' or the main ComfyUI folder."; + JOptionPane.showMessageDialog(parentDialog, msg, "Installation Error", JOptionPane.ERROR_MESSAGE); + return; + } + + // --- IMPROVEMENT 1: CLEANUP LEGACY/CONFLICTING FILES --- + String[] conflictingFiles = {"comfyui_to_downloader.py", "comfyui-model-downloader.py"}; + for (String conflict : conflictingFiles) { + File conflictFile = new File(customNodesDir, conflict); + if (conflictFile.exists()) { + System.out.println("Cleaning up legacy file: " + conflictFile.getAbsolutePath()); + conflictFile.delete(); + } + } + + configService.setComfyUIPath(comfyPath); + + File targetLink = new File(customNodesDir, "comfyui-model-downloader"); + File sourceDir = new File(System.getProperty("user.dir"), "comfyui-model-downloader"); + + // --- IMPROVEMENT 2: SOURCE VERIFICATION --- + if (!sourceDir.exists()) { + File parentSource = new File(new File(System.getProperty("user.dir")).getParent(), "comfyui-model-downloader"); + if (parentSource.exists()) sourceDir = parentSource; + } + + if (!sourceDir.exists()) { + File altSource = new File(System.getProperty("user.dir")); + if (new File(altSource, "__init__.py").exists() && new File(altSource, "web").exists()) { + sourceDir = altSource; + } + } + + if (!sourceDir.exists() || !new File(sourceDir, "__init__.py").exists()) { + JOptionPane.showMessageDialog(parentDialog, "Source extension folder is invalid or incomplete!\nLooked in: " + sourceDir.getAbsolutePath(), "Error", JOptionPane.ERROR_MESSAGE); + return; + } + + try { + // --- IMPROVEMENT 3: ROBUST DELETION OF OLD DIRECTORY/LINK --- + if (targetLink.exists()) { + if (System.getProperty("os.name").toLowerCase().contains("win")) { + Runtime.getRuntime().exec("cmd /c rmdir \"" + targetLink.getAbsolutePath() + "\"").waitFor(); + if (targetLink.exists()) Runtime.getRuntime().exec("cmd /c del /F /Q \"" + targetLink.getAbsolutePath() + "\"").waitFor(); + } + if (targetLink.exists()) deleteDirectory(targetLink); + } + + boolean success = false; + String method = "Unknown"; + + if (System.getProperty("os.name").toLowerCase().contains("win")) { + String cmdJ = String.format("cmd /c mklink /J \"%s\" \"%s\"", targetLink.getAbsolutePath(), sourceDir.getAbsolutePath()); + if (Runtime.getRuntime().exec(cmdJ).waitFor() == 0) { + success = true; method = "Junction"; + } else { + String cmdS = String.format("cmd /c mklink /D \"%s\" \"%s\"", targetLink.getAbsolutePath(), sourceDir.getAbsolutePath()); + if (Runtime.getRuntime().exec(cmdS).waitFor() == 0) { + success = true; method = "Symlink"; + } else { + copyDirectory(sourceDir, targetLink); + success = targetLink.exists(); method = "Copy"; + } + } + } else { + java.nio.file.Files.createSymbolicLink(targetLink.toPath(), sourceDir.toPath()); + success = true; method = "Symlink"; + } + + if (success) { + writeExtensionConfig(targetLink.exists() && method.equals("Copy") ? targetLink : sourceDir); + if (new File(targetLink, "__init__.py").exists()) { + String msg = "🚀 ComfyUI Bridge installed successfully!\n\n" + + "Method: " + method + "\n" + + "Location: " + targetLink.getAbsolutePath() + "\n\n" + + "Please RESTART ComfyUI now to see the rocket icon."; + JOptionPane.showMessageDialog(parentDialog, msg, "Success", JOptionPane.INFORMATION_MESSAGE); + parentDialog.dispose(); + } else { + throw new Exception("Installation successful but __init__.py not found at target!"); + } + } else { + throw new Exception("All installation methods failed."); + } + } catch (Exception e) { + JOptionPane.showMessageDialog(parentDialog, "Critical failure during installation: " + e.getMessage(), "Error", JOptionPane.ERROR_MESSAGE); + } + } + + private File findCustomNodesDeep(File dir, int depth) { + if (depth > 2) return null; // Limit depth to avoid performance issues + File[] files = dir.listFiles(); + if (files == null) return null; + + // Check immediate children first + for (File f : files) { + if (f.isDirectory() && f.getName().equalsIgnoreCase("custom_nodes")) { + return f; + } + } + + // Recurse + for (File f : files) { + if (f.isDirectory() && !f.getName().startsWith(".") && !f.getName().equalsIgnoreCase("node_modules")) { + File found = findCustomNodesDeep(f, depth + 1); + if (found != null) return found; + } + } + return null; + } + + private void writeExtensionConfig(File dir) { + try { + File webDir = new File(dir, "web"); + if (!webDir.exists()) webDir.mkdirs(); + JSONObject config = new JSONObject(); + config.put("token", configService.getApiToken()); + Files.writeString(new File(webDir, "config.json").toPath(), config.toString(4)); + } catch (Exception e) { e.printStackTrace(); } + } + + private void deleteDirectory(File dir) { + File[] files = dir.listFiles(); + if (files != null) for (File f : files) deleteDirectory(f); + dir.delete(); + } + + private void copyDirectory(File s, File d) throws IOException { + if (s.isDirectory()) { + if (!d.exists()) d.mkdirs(); + String[] children = s.list(); + if (children != null) for (String c : children) copyDirectory(new File(s, c), new File(d, c)); + } else Files.copy(s.toPath(), d.toPath(), StandardCopyOption.REPLACE_EXISTING); } private void loadFile(File file) { try { currentFileName = file.getName(); - String content = workflowService.extractWorkflow(file); - jsonInputArea.setText(content); + jsonInputArea.setText(workflowService.extractWorkflow(file)); analyzeJsonContent(); } catch (IOException ex) { JOptionPane.showMessageDialog(this, "Error: " + ex.getMessage()); } } private void analyzeJsonContent() { String text = jsonInputArea.getText(); - if (text.isEmpty()) return; + if (text == null || text.isEmpty()) return; modelsToDownload = analyzer.analyze(text, currentFileName); tableModel.setRowCount(0); - String baseModelsPath = modelsPathField.getText(); - + String base = modelsPathField.getText(); for (int i = 0; i < modelsToDownload.size(); i++) { ModelInfo info = modelsToDownload.get(i); String type = info.getType() != null ? info.getType() : "checkpoints"; - String targetSubDir = "models" + File.separator + type; - - // Check if file already exists locally - String status = info.getUrl().equals("MISSING") ? "Idle" : "✅ Known Good"; - boolean isSelected = true; - Path localPath = Paths.get(baseModelsPath, info.getSave_path() != null ? info.getSave_path() : type, info.getName()); - boolean existsLocally = Files.exists(localPath); - - if (existsLocally) { - status = "✅ Already exists"; - isSelected = false; // Deselect if already present - try { - long localSizeBytes = Files.size(localPath); - info.setSize(searchService.formatSize(localSizeBytes)); - } catch (IOException e) { - info.setSize("Local (Error)"); - } - } - - tableModel.addRow(new Object[]{isSelected, info.getType(), info.getName(), info.getSize(), info.getPopularity(), targetSubDir, info.getUrl(), status}); - - // Only fetch remote size if NOT already existing locally - if (!existsLocally && !info.getUrl().equals("MISSING")) { - fetchSizeInBackground(info, i); - } + Path local = Paths.get(base, info.getSave_path() != null ? info.getSave_path() : type, info.getName()); + boolean exists = Files.exists(local); + String status = exists ? "✅ Already exists" : (info.getUrl().equals("MISSING") ? "Idle" : "✅ Known Good"); + tableModel.addRow(new Object[]{!exists, info.getType(), info.getName(), info.getSize(), info.getPopularity(), "models/" + type, info.getUrl(), status}); } downloadButton.setEnabled(!modelsToDownload.isEmpty()); } private void searchMissingOnline() { if (modelsToDownload == null) return; - int rowCount = tableModel.getRowCount(); - boolean[] selectedForSearch = new boolean[rowCount]; - boolean anyMissingSelected = false; - for (int i = 0; i < rowCount; i++) { - 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 (!anyMissingSelected) { - JOptionPane.showMessageDialog(this, "No missing models selected for Deep Search."); - return; - } - statusLabel.setText("Deep Search in progress..."); - searchService.searchOnline(modelsToDownload, selectedForSearch, jsonInputArea.getText(), currentFileName, + statusLabel.setText("Searching..."); + boolean[] selected = new boolean[tableModel.getRowCount()]; + for (int i = 0; i < selected.length; i++) selected[i] = (Boolean) tableModel.getValueAt(i, 0); + searchService.searchOnline(modelsToDownload, selected, 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(targetSubDir, idx, 5); tableModel.setValueAt(info.getUrl(), idx, 6); tableModel.setValueAt("✅ Found", idx, 7); }), - () -> SwingUtilities.invokeLater(() -> statusLabel.setText("Deep Search finished.")) + () -> SwingUtilities.invokeLater(() -> statusLabel.setText("Search finished.")) ); } @@ -665,7 +772,7 @@ public class Main extends JFrame { try { List errors = new ArrayList<>(); Map> hashToPaths = new HashMap<>(); - + List allFiles = Files.walk(root.toPath()) .filter(Files::isRegularFile) .filter(p -> { @@ -679,7 +786,7 @@ public class Main extends JFrame { Path p = allFiles.get(i); final int current = i + 1; SwingUtilities.invokeLater(() -> statusLabel.setText("Checking (" + current + "/" + total + "): " + p.getFileName())); - + IModelValidator.ValidationResult res = modelValidator.validateFile(p.toFile()); if (!res.ok) { errors.add(res); @@ -697,158 +804,114 @@ public class Main extends JFrame { SwingUtilities.invokeLater(() -> { statusLabel.setText("Scan finished. Found " + errors.size() + " issues and " + duplicates.size() + " duplicate sets."); - + if (errors.isEmpty() && duplicates.isEmpty()) { JOptionPane.showMessageDialog(this, "All " + total + " models verified successfully!", "Verification Complete", JOptionPane.INFORMATION_MESSAGE); return; } - StringBuilder sb = new StringBuilder(); if (!errors.isEmpty()) { - sb.append("⚠️ CORRUPTED FILES (").append(errors.size()).append("):\n"); + StringBuilder sb = new StringBuilder("The following " + errors.size() + " files appear to be corrupted or invalid:\n\n"); for (IModelValidator.ValidationResult err : errors) { - sb.append("- ").append(new File(err.filePath).getName()).append(" (").append(err.message).append(")\n"); + sb.append("- ").append(new File(err.filePath).getName()) + .append(" (").append(err.message).append(")\n") + .append(" Path: ").append(err.filePath).append("\n\n"); } - sb.append("\n"); - } - if (!duplicates.isEmpty()) { - sb.append("👯 DUPLICATES FOUND (").append(duplicates.size()).append(" sets):\n"); - for (Map.Entry> entry : duplicates.entrySet()) { - sb.append("Hash: ").append(entry.getKey().substring(0, 12)).append("...\n"); - for (Path p : entry.getValue()) { - sb.append(" - ").append(p.toFile().getAbsolutePath()).append(" (").append(formatSize(p.toFile().length())).append(")\n"); + JTextArea textArea = new JTextArea(sb.toString()); + textArea.setEditable(false); + JScrollPane scrollPane = new JScrollPane(textArea); + scrollPane.setPreferredSize(new Dimension(800, 500)); + + Object[] options = {"OK", "Delete All Corrupted Files"}; + int choice = JOptionPane.showOptionDialog(this, scrollPane, "Verification Results - Issues Found", + JOptionPane.YES_NO_OPTION, JOptionPane.WARNING_MESSAGE, null, options, options[0]); + + if (choice == 1) { // Delete All Corrupted Files + int confirm = JOptionPane.showConfirmDialog(this, + "Are you sure you want to delete these " + errors.size() + " files?\nThis action cannot be undone.", + "Confirm Deletion", JOptionPane.YES_NO_OPTION, JOptionPane.WARNING_MESSAGE); + + if (confirm == JOptionPane.YES_OPTION) { + int deletedCount = 0; + for (IModelValidator.ValidationResult err : errors) { + File f = new File(err.filePath); + if (f.exists() && f.delete()) { + deletedCount++; + } + } + JOptionPane.showMessageDialog(this, "Deleted " + deletedCount + " files."); + statusLabel.setText("Cleanup finished. Deleted " + deletedCount + " files."); } - sb.append("\n"); } } - JTextArea textArea = new JTextArea(sb.toString()); - textArea.setEditable(false); - JScrollPane scrollPane = new JScrollPane(textArea); - scrollPane.setPreferredSize(new Dimension(850, 550)); - - Object[] options = {"OK", "Delete All Corrupted", "Manage Duplicates..."}; - int choice = JOptionPane.showOptionDialog(this, scrollPane, "Verification & Duplicate Results", - JOptionPane.DEFAULT_OPTION, JOptionPane.WARNING_MESSAGE, null, options, options[0]); - - if (choice == 1) { // Delete All Corrupted - handleDeleteCorrupted(errors); - } else if (choice == 2) { // Manage Duplicates - handleManageDuplicates(duplicates); + if (!duplicates.isEmpty()) { + showDuplicatesDialog(duplicates); } }); } catch (IOException e) { - SwingUtilities.invokeLater(() -> JOptionPane.showMessageDialog(this, "Error: " + e.getMessage())); + SwingUtilities.invokeLater(() -> JOptionPane.showMessageDialog(this, "Error during verification: " + e.getMessage())); } }).start(); } - private void handleDeleteCorrupted(List errors) { - int confirm = JOptionPane.showConfirmDialog(this, "Delete " + errors.size() + " corrupted files?", "Confirm", JOptionPane.YES_NO_OPTION); - if (confirm == JOptionPane.YES_OPTION) { - int deleted = 0; - for (IModelValidator.ValidationResult err : errors) { - File f = new File(err.filePath); - if (f.exists() && f.delete()) { - hashRegistry.unregister(f); - deleted++; - } + private void showDuplicatesDialog(Map> duplicates) { + StringBuilder sb = new StringBuilder("Storage Optimizer - Duplicate Models Found:\n\n"); + sb.append("The following files have identical content (SHA-256 match).\n"); + sb.append("You might want to delete redundant copies to save space.\n\n"); + + for (Map.Entry> entry : duplicates.entrySet()) { + sb.append("SHA-256: ").append(entry.getKey()).append("\n"); + for (Path p : entry.getValue()) { + sb.append(" -> ").append(p.toAbsolutePath()).append("\n"); } - JOptionPane.showMessageDialog(this, "Deleted " + deleted + " files."); + sb.append("\n"); } + + JTextArea textArea = new JTextArea(sb.toString()); + textArea.setEditable(false); + JScrollPane scrollPane = new JScrollPane(textArea); + scrollPane.setPreferredSize(new Dimension(900, 600)); + JOptionPane.showMessageDialog(this, scrollPane, "Duplicates Found", JOptionPane.INFORMATION_MESSAGE); } - private void handleManageDuplicates(Map> duplicates) { - long totalSavings = 0; - for (List paths : duplicates.values()) { - for (int i = 1; i < paths.size(); i++) totalSavings += paths.get(i).toFile().length(); - } - - int confirm = JOptionPane.showConfirmDialog(this, - "Keep the first file of each set and delete others?\nEstimated savings: " + formatSize(totalSavings), - "Storage Optimizer", JOptionPane.YES_NO_OPTION); - - if (confirm == JOptionPane.YES_OPTION) { - int deletedCount = 0; - for (List paths : duplicates.values()) { - for (int i = 1; i < paths.size(); i++) { - File f = paths.get(i).toFile(); - if (f.exists() && f.delete()) { - hashRegistry.unregister(f); - deletedCount++; - } - } + private void updateDownloadManagerSelection() { + if (downloadManager == null || tableModel == null) return; + boolean[] selected = new boolean[tableModel.getRowCount()]; + for (int i = 0; i < selected.length; 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) {} + public void dragOver(DropTargetDragEvent dtde) {} + public void dropActionChanged(DropTargetDragEvent dtde) {} + public void dragExit(DropTargetEvent dte) {} + public void drop(DropTargetDropEvent dtde) { + try { + dtde.acceptDrop(DnDConstants.ACTION_COPY); + java.util.List files = (java.util.List) dtde.getTransferable().getTransferData(DataFlavor.javaFileListFlavor); + if (!files.isEmpty()) loadFile(files.get(0)); + } catch (Exception e) {} } - JOptionPane.showMessageDialog(this, "Freed " + formatSize(totalSavings) + "."); - } + }); + } + + private void performSystemShutdown() { + try { Runtime.getRuntime().exec("shutdown /s /t 60"); } catch (IOException ignored) {} } private void showHelpDialog() { - String helpText = "" + - "

ComfyUI Model Downloader Help

" + - "

This tool automates the process of finding and downloading models required for your ComfyUI workflows.

" + - - "

🚀 Core Features

" + - "
    " + - "
  • Smart Extraction: Drag and drop JSON or PNG workflow files. The tool extracts models even from PNG metadata.
    • " + - "
  • Deep Search: Automatically finds URLs on Hugging Face (Official & Community) and Civitai.
    • " + - "
  • AI Scouting: Uses Gemini AI to understand the workflow context and find the exact right model versions.
    • " + - "
  • Auto-Organization: Automatically sorts models into correct subfolders (checkpoints, loras, vae, controlnet, etc.).
    • " + - "
  • Existing File Detection: Detects models you already have, calculates their local size, and prevents redundant downloads.
    • " + - "
" + - - "

🛡️ Quality & Security

" + - "
    " + - "
  • Model Verification: Use the Verify button to scan your local models for corrupted files (e.g., interrupted downloads or broken Safetensors).
    • " + - "
  • Encrypted Vault: Your optional API keys are stored in settings.vault using AES-256 encryption, protected by your master password.
    • " + - "
  • Privacy First: No data is ever sent to third-party servers except for official API requests to Google, Hugging Face, or Civitai.
    • " + - "
" + - - "

⚙️ Advanced Settings

" + - "
    " + - "
  • Background Mode: Close the window to keep the downloader running in the System Tray.
    • " + - "
  • Auto-Shutdown: Automatically shut down your PC after the download queue is finished.
    • " + - "
  • Dark Mode: Modern dark interface for better workflow integration.
    • " + - "
  • Model List Import: Import your own JSON/CSV lists to prioritize your trusted sources.
    • " + - "
" + - - "

🔑 API Keys (Optional)

" + - "
    " + - "
  • Gemini Key: Recommended for the best 'AI Scouting' results and repository discovery.
    • " + - "
  • HF Token: Required to download Gated Models (like FLUX.1-dev or certain SD3 versions).
    • " + - "
" + - ""; - - 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()); - String s = formatSize(size); - info.setSize(s); - SwingUtilities.invokeLater(() -> tableModel.setValueAt(s, rowIndex, 3)); - }).start(); - } - - private String formatSize(long bytes) { - if (bytes <= 0) return "Unknown"; - if (bytes < 1024 * 1024) return (bytes / 1024) + " KB"; - if (bytes < 1024 * 1024 * 1024) return String.format("%.1f MB", bytes / (1024.0 * 1024.0)); - return String.format("%.2f GB", bytes / (1024.0 * 1024.0 * 1024.0)); + JOptionPane.showMessageDialog(this, "ComfyUI Model Downloader\n\n1. Set your Models path\n2. Load or Drag&Drop a Workflow\n3. Click Start Queue", "Help", JOptionPane.INFORMATION_MESSAGE); } public static void main(String[] args) { - // Enable modern window decorations BEFORE anything else - com.formdev.flatlaf.FlatLaf.setUseNativeWindowDecorations(true); - + FlatLaf.setUseNativeWindowDecorations(true); AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(AppConfig.class); - Main main = context.getBean(Main.class); - main.launch(args); + context.getBean(Main.class).launch(args); } - @Configuration - @ComponentScan("de.tki.comfymodels") - public static class AppConfig {} -} \ No newline at end of file + @Configuration @ComponentScan("de.tki.comfymodels") public static class AppConfig {} +} 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 5191631..c4d8d0e 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java @@ -122,6 +122,19 @@ public class ConfigService { return null; } + public String getComfyUIPath() { return settings.optString("comfyui_path", ""); } + public void setComfyUIPath(String path) { settings.put("comfyui_path", path); save(); } + + public String getApiToken() { + String token = settings.optString("api_token", ""); + if (token.isEmpty() && isUnlocked()) { + token = java.util.UUID.randomUUID().toString().replace("-", ""); + settings.put("api_token", token); + save(); + } + return token; + } + public boolean isUnlocked() { return masterPassword != null; } 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 efdf690..1164879 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java +++ b/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java @@ -29,9 +29,14 @@ public class DefaultDownloadManager implements IDownloadManager { private ConfigService configService; private void safeUpdateStatus(int index, String status, BiConsumer statusUpdater) { - if (completedIndices.contains(index)) return; + // Only mark as completed if it's a final state + boolean isFinal = status.contains("✅ Finished") || + status.contains("Stopped") || + status.contains("Skipped") || + status.contains("Already exists") || + status.startsWith("❌ Error"); - if (status.contains("✅ Finished") || status.contains("Stopped") || status.contains("Skipped") || status.contains("Already exists")) { + if (isFinal) { completedIndices.add(index); } statusUpdater.accept(index, status); @@ -43,39 +48,32 @@ public class DefaultDownloadManager implements IDownloadManager { isPaused = false; this.currentSelection = selectedIndices; completedIndices.clear(); + + java.util.concurrent.atomic.AtomicBoolean finishedCalled = new java.util.concurrent.atomic.AtomicBoolean(false); + 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; - } - } - + int size = models.size(); + java.util.concurrent.CompletableFuture[] futures = new java.util.concurrent.CompletableFuture[size]; + + for (int i = 0; i < size; i++) { final int index = i; - if (!isSelected(index)) { - safeUpdateStatus(index, "Skipped (Not Selected)", statusUpdater); - futures[i] = java.util.concurrent.CompletableFuture.completedFuture(null); + if (isStopped) { + futures[index] = java.util.concurrent.CompletableFuture.completedFuture(null); continue; } - ModelInfo info = models.get(i); - futures[i] = java.util.concurrent.CompletableFuture.runAsync(() -> { - 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; } + if (!isSelected(index)) { + safeUpdateStatus(index, "Skipped (Not Selected)", statusUpdater); + futures[index] = java.util.concurrent.CompletableFuture.completedFuture(null); + continue; + } + + ModelInfo info = models.get(index); + futures[index] = java.util.concurrent.CompletableFuture.runAsync(() -> { + if (isStopped || !isSelected(index)) return; + try { String subPath = info.getSave_path(); if (subPath == null || subPath.isEmpty() || "default".equalsIgnoreCase(subPath)) { @@ -87,17 +85,24 @@ public class DefaultDownloadManager implements IDownloadManager { Files.createDirectories(targetDir); downloadWithResume(info, targetDir.resolve(info.getName()), index, statusUpdater); } catch (Exception e) { - if (isStopped || !isSelected(index)) { - safeUpdateStatus(index, !isSelected(index) ? "Skipped (Unchecked)" : "Stopped", statusUpdater); - } else { + if (!isStopped && isSelected(index)) { String msg = e.getMessage() != null ? e.getMessage() : e.getClass().getSimpleName(); - safeUpdateStatus(index, "Error: " + msg, statusUpdater); + safeUpdateStatus(index, "❌ Error: " + msg, statusUpdater); } } }, executor); } - java.util.concurrent.CompletableFuture.allOf(futures).join(); - } finally { onFinished.run(); } + + // Wait for all tasks to complete (including skipped ones) + try { + java.util.concurrent.CompletableFuture.allOf(futures).join(); + } catch (Exception ignored) {} + + } finally { + if (finishedCalled.compareAndSet(false, true)) { + onFinished.run(); + } + } }).start(); } diff --git a/src/main/java/de/tki/comfymodels/service/impl/RestBridgeService.java b/src/main/java/de/tki/comfymodels/service/impl/RestBridgeService.java new file mode 100644 index 0000000..f6e0b3f --- /dev/null +++ b/src/main/java/de/tki/comfymodels/service/impl/RestBridgeService.java @@ -0,0 +1,125 @@ +package de.tki.comfymodels.service.impl; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; +import org.json.JSONObject; +import org.springframework.stereotype.Service; + +import jakarta.annotation.PostConstruct; +import jakarta.annotation.PreDestroy; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.util.function.Consumer; + +@Service +public class RestBridgeService { + + private HttpServer server; + private Consumer workflowConsumer; + private int port = 12345; + private String expectedApiToken; + + public void setPort(int port) { + this.port = port; + } + + public void setApiToken(String token) { + this.expectedApiToken = token; + } + + public void setWorkflowConsumer(Consumer consumer) { + this.workflowConsumer = consumer; + } + + public void startServer() { + if (server != null) return; + try { + server = HttpServer.create(new InetSocketAddress("127.0.0.1", port), 0); + server.createContext("/import", new ImportHandler()); + server.setExecutor(null); + server.start(); + System.out.println("REST Bridge started on port " + port); + } catch (IOException e) { + System.err.println("Failed to start REST Bridge: " + e.getMessage()); + } + } + + @PreDestroy + public void stopServer() { + if (server != null) { + server.stop(0); + server = null; + } + } + + private class ImportHandler implements HttpHandler { + @Override + public void handle(HttpExchange exchange) throws IOException { + try { + // Add CORS headers to EVERY response + exchange.getResponseHeaders().add("Access-Control-Allow-Origin", "*"); + exchange.getResponseHeaders().add("Access-Control-Allow-Methods", "POST, OPTIONS, GET"); + exchange.getResponseHeaders().add("Access-Control-Allow-Headers", "Content-Type, Authorization"); + + // Handle preflight OPTIONS request + if ("OPTIONS".equalsIgnoreCase(exchange.getRequestMethod())) { + exchange.sendResponseHeaders(204, -1); + return; + } + + // SECURITY: Validate API Token + String authHeader = exchange.getRequestHeaders().getFirst("Authorization"); + if (expectedApiToken != null && !expectedApiToken.isEmpty()) { + if (authHeader == null || !authHeader.equals("Bearer " + expectedApiToken)) { + System.err.println("REST Bridge: Unauthorized request (Invalid or Missing Token)"); + exchange.sendResponseHeaders(401, -1); + return; + } + } + + if (!"POST".equalsIgnoreCase(exchange.getRequestMethod())) { + exchange.sendResponseHeaders(405, -1); + return; + } + + try (InputStream is = exchange.getRequestBody()) { + String body = new String(is.readAllBytes(), StandardCharsets.UTF_8); + + if (workflowConsumer != null) { + workflowConsumer.accept(body); + + String response = "{\"status\": \"success\", \"message\": \"Workflow received\"}"; + byte[] responseBytes = response.getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().set("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, responseBytes.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(responseBytes); + } + } else { + String response = "{\"status\": \"error\", \"message\": \"App not ready\"}"; + byte[] responseBytes = response.getBytes(StandardCharsets.UTF_8); + exchange.sendResponseHeaders(503, responseBytes.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(responseBytes); + } + } + } + } catch (Exception e) { + String response = "{\"status\": \"error\", \"message\": \"" + e.getMessage() + "\"}"; + byte[] responseBytes = response.getBytes(StandardCharsets.UTF_8); + try { + exchange.sendResponseHeaders(500, responseBytes.length); + try (OutputStream os = exchange.getResponseBody()) { + os.write(responseBytes); + } + } catch (IOException ignored) {} + } finally { + exchange.close(); + } + } + } +} diff --git a/src/test/java/de/tki/comfymodels/service/InstallationPathTest.java b/src/test/java/de/tki/comfymodels/service/InstallationPathTest.java new file mode 100644 index 0000000..df7c33d --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/InstallationPathTest.java @@ -0,0 +1,88 @@ +package de.tki.comfymodels.service; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; + +import static org.junit.jupiter.api.Assertions.*; + +class InstallationPathTest { + + @TempDir + Path tempDir; + + @Test + void testStandardLayout() throws IOException { + Path customNodes = tempDir.resolve("custom_nodes"); + Files.createDirectory(customNodes); + + File inputPath = tempDir.toFile(); + File customNodesDir = findCustomNodes(inputPath); + + assertNotNull(customNodesDir); + assertEquals(customNodes.toAbsolutePath().toString(), customNodesDir.getAbsolutePath()); + } + + @Test + void testPortableLayout() throws IOException { + Path comfyDir = tempDir.resolve("ComfyUI"); + Files.createDirectory(comfyDir); + Path customNodes = comfyDir.resolve("custom_nodes"); + Files.createDirectory(customNodes); + + File inputPath = tempDir.toFile(); + File customNodesDir = findCustomNodes(inputPath); + + assertNotNull(customNodesDir); + assertEquals(customNodes.toAbsolutePath().toString(), customNodesDir.getAbsolutePath()); + } + + @Test + void testDesktopAppLayout() throws IOException { + Path resources = tempDir.resolve("resources"); + Files.createDirectory(resources); + Path comfyDir = resources.resolve("ComfyUI"); + Files.createDirectory(comfyDir); + Path customNodes = comfyDir.resolve("custom_nodes"); + Files.createDirectory(customNodes); + + File inputPath = tempDir.toFile(); + File customNodesDir = findCustomNodes(inputPath); + + assertNotNull(customNodesDir); + assertEquals(customNodes.toAbsolutePath().toString(), customNodesDir.getAbsolutePath()); + } + + @Test + void testDirectSelection() throws IOException { + Path customNodes = tempDir.resolve("custom_nodes"); + Files.createDirectory(customNodes); + + File inputPath = customNodes.toFile(); + File customNodesDir = findCustomNodes(inputPath); + + assertNotNull(customNodesDir); + assertEquals(customNodes.toAbsolutePath().toString(), customNodesDir.getAbsolutePath()); + } + + /** + * Logic extracted from Main.java for testing + */ + private File findCustomNodes(File inputPath) { + File c1 = new File(inputPath, "custom_nodes"); + File c2 = new File(inputPath, "ComfyUI" + File.separator + "custom_nodes"); + File c3 = new File(inputPath, "resources" + File.separator + "ComfyUI" + File.separator + "custom_nodes"); + File c4 = inputPath; + + if (c1.exists() && c1.isDirectory()) return c1; + if (c2.exists() && c2.isDirectory()) return c2; + if (c3.exists() && c3.isDirectory()) return c3; + if (c4.getName().equalsIgnoreCase("custom_nodes") && c4.isDirectory()) return c4; + + return null; + } +} diff --git a/src/test/java/de/tki/comfymodels/service/RestBridgeServiceTest.java b/src/test/java/de/tki/comfymodels/service/RestBridgeServiceTest.java new file mode 100644 index 0000000..5b0ca5b --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/RestBridgeServiceTest.java @@ -0,0 +1,53 @@ +package de.tki.comfymodels.service; + +import de.tki.comfymodels.service.impl.RestBridgeService; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class RestBridgeServiceTest { + + private RestBridgeService restBridge; + private final int testPort = 12346; + private final String testToken = "test-token"; + + @BeforeEach + void setUp() { + restBridge = new RestBridgeService(); + restBridge.setPort(testPort); + restBridge.setApiToken(testToken); + restBridge.startServer(); + } + + @AfterEach + void tearDown() { + restBridge.stopServer(); + } + + @Test + void testWorkflowImport() throws Exception { + AtomicReference receivedWorkflow = new AtomicReference<>(); + restBridge.setWorkflowConsumer(receivedWorkflow::set); + + String testJson = "{\"nodes\": []}"; + HttpClient client = HttpClient.newHttpClient(); + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create("http://127.0.0.1:" + testPort + "/import")) + .header("Authorization", "Bearer " + testToken) + .POST(HttpRequest.BodyPublishers.ofString(testJson)) + .build(); + + HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); + + assertEquals(200, response.statusCode()); + assertEquals(testJson, receivedWorkflow.get()); + } +}