diff --git a/.gitignore b/.gitignore index abb1020..5f7f37c 100644 --- a/.gitignore +++ b/.gitignore @@ -16,15 +16,19 @@ Thumbs.db .DS_Store # --- Project Specific --- -# Sensitive API keys +# Sensitive API keys and secrets app_settings.json settings.vault +boot.key +.env # User Model List (internal DB) uploaded_models.json +pending_downloads.json # Local config & assets workflow.json *.png *.bat +*.zip !src/main/resources/known_good_urls.json diff --git a/README.md b/README.md index c80d1d5..1656c8c 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,26 @@ A robust, feature-rich Java Swing application designed to automate the process o java -jar target/comfymodeldownloader-1.0-SNAPSHOT.jar ``` +## ๐Ÿ First Start & Model List Setup + +Upon your first launch, it is highly recommended to import a master model list. This list helps the application immediately recognize models and their correct download sources. + +1. **Prepare your JSON**: Create or obtain a `model-list.json` file. It should follow this structure: + ```json + { + "models": [ + { + "name": "SDXL Base 1.0", + "filename": "sd_xl_base_1.0.safetensors", + "type": "checkpoints", + "url": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors" + } + ] + } + ``` +2. **Import**: Click the **Import Model List (JSON)...** button in the workflow panel. +3. **Verification**: Once imported, the application will use this list as a primary reference for all future workflow analyses. + ## ๐Ÿ“– How to Use 1. **Setup Path**: Enter or browse to your local ComfyUI models directory. diff --git a/pom.xml b/pom.xml index a1d554f..dfb25f2 100644 --- a/pom.xml +++ b/pom.xml @@ -49,13 +49,14 @@ + ComfyUIModelDownloader org.apache.maven.plugins maven-compiler-plugin 3.11.0 - 11 + 17 @@ -69,11 +70,27 @@ shade + false de.tki.comfymodels.Main + + + *:* + + META-INF/*.SF + META-INF/*.DSA + META-INF/*.RSA + *.vault + *.key + app_settings.json + pending_downloads.json + workflow.json + + + diff --git a/run_downloader.bat b/run_downloader.bat deleted file mode 100644 index b5f814b..0000000 --- a/run_downloader.bat +++ /dev/null @@ -1,13 +0,0 @@ -@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 3234001..118afeb 100644 --- a/src/main/java/de/tki/comfymodels/Main.java +++ b/src/main/java/de/tki/comfymodels/Main.java @@ -4,10 +4,13 @@ import de.tki.comfymodels.domain.ModelInfo; import de.tki.comfymodels.service.IDownloadManager; import de.tki.comfymodels.service.IModelAnalyzer; import de.tki.comfymodels.service.IWorkflowService; +import de.tki.comfymodels.service.IModelValidator; import de.tki.comfymodels.service.IModelSearchService; import de.tki.comfymodels.service.impl.ConfigService; import de.tki.comfymodels.service.impl.GeminiAIService; import de.tki.comfymodels.service.impl.ModelListService; +import org.json.JSONArray; +import org.json.JSONObject; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.annotation.AnnotationConfigApplicationContext; import org.springframework.context.annotation.ComponentScan; @@ -23,10 +26,16 @@ import java.awt.datatransfer.DataFlavor; import java.awt.dnd.*; import java.awt.event.MouseAdapter; import java.awt.event.MouseEvent; +import java.awt.event.ActionListener; +import java.awt.image.BufferedImage; import java.io.*; import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; import java.util.List; import java.util.ArrayList; +import java.util.stream.Collectors; @Component public class Main extends JFrame { @@ -34,6 +43,7 @@ public class Main extends JFrame { private final IDownloadManager downloadManager; private final IWorkflowService workflowService; private final IModelSearchService searchService; + private final IModelValidator modelValidator; @Autowired private ConfigService configService; @@ -47,6 +57,8 @@ public class Main extends JFrame { private JTextField modelsPathField; private JPasswordField geminiKeyField; private JPasswordField hfTokenField; + private JCheckBox backgroundCheck; + private JCheckBox shutdownCheck; private JLabel activeAiModelLabel; private JTextArea jsonInputArea; private DefaultTableModel tableModel; @@ -56,14 +68,24 @@ public class Main extends JFrame { private String currentFileName = "input.json"; public Main(IModelAnalyzer analyzer, IDownloadManager downloadManager, - IWorkflowService workflowService, IModelSearchService searchService) { + IWorkflowService workflowService, IModelSearchService searchService, + IModelValidator modelValidator) { this.analyzer = analyzer; this.downloadManager = downloadManager; this.workflowService = workflowService; this.searchService = searchService; + this.modelValidator = modelValidator; } - public void launch() { + public void launch(String[] args) { + boolean backgroundArg = false; + for (String arg : args) { + if ("--background".equalsIgnoreCase(arg)) { + backgroundArg = true; + break; + } + } + try { UIManager.setLookAndFeel(UIManager.getSystemLookAndFeelClassName()); } catch (Exception e) {} @@ -74,15 +96,20 @@ public class Main extends JFrame { SwingUtilities.invokeLater(() -> { initUI(); + setupTrayIcon(); loadSettingsIntoUI(); updateAiModelDisplay(); - // Add WindowListener to stop downloads on close + // Add WindowListener to handle background mode addWindowListener(new java.awt.event.WindowAdapter() { @Override public void windowClosing(java.awt.event.WindowEvent e) { - downloadManager.stop(); - System.exit(0); + if (configService.isBackgroundModeEnabled()) { + setVisible(false); + } else { + downloadManager.stop(); + System.exit(0); + } } }); @@ -90,6 +117,64 @@ public class Main extends JFrame { }); } + private void setupTrayIcon() { + 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 + g2.fillRect(0, 0, 16, 16); + g2.dispose(); + + PopupMenu popup = new PopupMenu(); + MenuItem showItem = new MenuItem("Show UI"); + showItem.addActionListener(e -> setVisible(true)); + MenuItem exitItem = new MenuItem("Exit"); + exitItem.addActionListener(e -> { + downloadManager.stop(); + System.exit(0); + }); + + popup.add(showItem); + popup.addSeparator(); + popup.add(exitItem); + + TrayIcon trayIcon = new TrayIcon(image, "ComfyUI Model Downloader", popup); + trayIcon.setImageAutoSize(true); + trayIcon.addActionListener(e -> setVisible(true)); + + try { + tray.add(trayIcon); + } catch (AWTException e) { + System.err.println("TrayIcon could not be added."); + } + } + + 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)); @@ -134,6 +219,12 @@ public class Main extends JFrame { 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()); + } } private void updateAiModelDisplay() { @@ -146,7 +237,7 @@ public class Main extends JFrame { private void initUI() { setTitle("ComfyUIModel-Downloader"); setSize(1450, 1000); - setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE); + setDefaultCloseOperation(JFrame.DO_NOTHING_ON_CLOSE); setLayout(new BorderLayout()); JPanel settingsPanel = new JPanel(new GridLayout(4, 1)); @@ -209,7 +300,20 @@ public class Main extends JFrame { JPanel infoRow = new JPanel(new FlowLayout(FlowLayout.LEFT)); JButton helpBtn = new JButton("โ„น Help & Security Information"); helpBtn.addActionListener(e -> showHelpDialog()); + + JButton verifyBtn = new JButton("๐Ÿ” Verify Local Models (Corruption Check)"); + verifyBtn.addActionListener(e -> verifyLocalModels()); + + backgroundCheck = new JCheckBox("Stay in Background on Close"); + backgroundCheck.addActionListener(e -> configService.setBackgroundModeEnabled(backgroundCheck.isSelected())); + + shutdownCheck = new JCheckBox("Shutdown after Queue"); + shutdownCheck.addActionListener(e -> configService.setShutdownAfterDownloadEnabled(shutdownCheck.isSelected())); + infoRow.add(helpBtn); + infoRow.add(verifyBtn); + infoRow.add(backgroundCheck); + infoRow.add(shutdownCheck); settingsPanel.add(pathRow); settingsPanel.add(apiRow); @@ -311,11 +415,22 @@ public class Main extends JFrame { } downloadButton.setEnabled(false); + updatePendingDownloadsPersistence(); downloadManager.startQueue(modelsToDownload, selected, modelsPathField.getText(), - (idx, status) -> SwingUtilities.invokeLater(() -> tableModel.setValueAt(status, idx, 7)), + (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(); + } }) ); }); @@ -350,6 +465,21 @@ public class Main extends JFrame { 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() { int rowCount = tableModel.getRowCount(); boolean[] selected = new boolean[rowCount]; @@ -389,10 +519,23 @@ public class Main extends JFrame { if (text.isEmpty()) return; modelsToDownload = analyzer.analyze(text, currentFileName); tableModel.setRowCount(0); + String baseModelsPath = modelsPathField.getText(); + for (int i = 0; i < modelsToDownload.size(); i++) { ModelInfo info = modelsToDownload.get(i); - 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"}); + 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()); + if (Files.exists(localPath)) { + status = "โœ… Already exists"; + isSelected = false; // Deselect if already present + } + + tableModel.addRow(new Object[]{isSelected, info.getType(), info.getName(), info.getSize(), info.getPopularity(), targetSubDir, info.getUrl(), status}); if (!info.getUrl().equals("MISSING")) fetchSizeInBackground(info, i); } downloadButton.setEnabled(!modelsToDownload.isEmpty()); @@ -432,6 +575,89 @@ public class Main extends JFrame { ); } + private void verifyLocalModels() { + String baseDir = modelsPathField.getText(); + if (baseDir == null || baseDir.isEmpty()) { + JOptionPane.showMessageDialog(this, "Please set a models directory first."); + return; + } + + File root = new File(baseDir); + if (!root.exists() || !root.isDirectory()) { + JOptionPane.showMessageDialog(this, "Invalid models directory."); + return; + } + + statusLabel.setText("Verifying models... please wait."); + new Thread(() -> { + try { + List errors = new ArrayList<>(); + List allFiles = Files.walk(root.toPath()) + .filter(Files::isRegularFile) + .filter(p -> { + String n = p.getFileName().toString().toLowerCase(); + return n.endsWith(".safetensors") || n.endsWith(".sft") || n.endsWith(".ckpt") || n.endsWith(".pth") || n.endsWith(".pt") || n.endsWith(".bin"); + }) + .collect(Collectors.toList()); + + int total = allFiles.size(); + for (int i = 0; i < total; i++) { + 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); + } + } + + SwingUtilities.invokeLater(() -> { + statusLabel.setText("Verification finished. Found " + errors.size() + " issues."); + if (errors.isEmpty()) { + JOptionPane.showMessageDialog(this, "All " + total + " models verified successfully!", "Verification Complete", JOptionPane.INFORMATION_MESSAGE); + } else { + 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") + .append(" Path: ").append(err.filePath).append("\n\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."); + } + } + } + }); + } catch (IOException e) { + SwingUtilities.invokeLater(() -> JOptionPane.showMessageDialog(this, "Error during verification: " + e.getMessage())); + } + }).start(); + } + private void showHelpDialog() { String helpText = "" + "

ComfyUIModel-Downloader Help

" + @@ -471,7 +697,7 @@ public class Main extends JFrame { public static void main(String[] args) { AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(AppConfig.class); Main main = context.getBean(Main.class); - main.launch(); + main.launch(args); } @Configuration diff --git a/src/main/java/de/tki/comfymodels/domain/ModelInfo.java b/src/main/java/de/tki/comfymodels/domain/ModelInfo.java index 6d86282..7de39dc 100644 --- a/src/main/java/de/tki/comfymodels/domain/ModelInfo.java +++ b/src/main/java/de/tki/comfymodels/domain/ModelInfo.java @@ -14,6 +14,8 @@ public class ModelInfo { private String reference; private String filename; + public ModelInfo() {} + public ModelInfo(String type, String name, String url) { this.type = type; this.name = name; @@ -23,6 +25,7 @@ public class ModelInfo { public String getType() { return type; } public void setType(String type) { this.type = type; } public String getName() { return name; } + public void setName(String name) { this.name = name; } public String getUrl() { return url; } public void setUrl(String url) { this.url = url; } public String getSize() { return size; } diff --git a/src/main/java/de/tki/comfymodels/service/IModelValidator.java b/src/main/java/de/tki/comfymodels/service/IModelValidator.java new file mode 100644 index 0000000..b9e7681 --- /dev/null +++ b/src/main/java/de/tki/comfymodels/service/IModelValidator.java @@ -0,0 +1,19 @@ +package de.tki.comfymodels.service; + +import java.io.File; + +public interface IModelValidator { + ValidationResult validateFile(File file); + + class ValidationResult { + public final boolean ok; + public final String message; + public final String filePath; + + public ValidationResult(boolean ok, String message, String filePath) { + this.ok = ok; + this.message = message; + this.filePath = filePath; + } + } +} 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 78ecb85..e39cbf3 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java @@ -11,10 +11,6 @@ import java.util.*; import java.util.regex.Matcher; import java.util.regex.Pattern; -import jakarta.annotation.PostConstruct; -import java.io.InputStream; -import java.nio.charset.StandardCharsets; - @Service public class ComfyModelAnalyzer implements IModelAnalyzer { @@ -33,31 +29,11 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { 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<>(); - - @PostConstruct - public void init() { - try (InputStream is = getClass().getResourceAsStream("/known_good_urls.json")) { - if (is != null) { - String json = new String(is.readAllBytes(), StandardCharsets.UTF_8); - JSONObject jo = new JSONObject(json); - for (String key : jo.keySet()) { - knownGoodUrls.put(key.toLowerCase(), jo.getString(key)); - } - } - } catch (Exception e) { - System.err.println("Failed to load known_good_urls.json: " + e.getMessage()); - } + static final Pattern NAKED_URL = Pattern.compile("(https?://[a-zA-Z0-9\\-\\.\\/\\?\\=\\&\\%]+(?:\\.(?:safetensors|sft|ckpt|pth|pt|bin|onnx|yaml))(?:\\?[^\\s\\\"\\']*)?)", Pattern.CASE_INSENSITIVE); } @Override public List analyze(String jsonText, String fileName) { - // Fallback falls Spring Injection fehlt (Test-Modus) - if (aiService == null) aiService = new LocalAIService(); - if (knownGoodUrls.isEmpty()) init(); // Fallback fรผr Tests ohne Spring Context - List results = new ArrayList<>(); if (jsonText == null || jsonText.trim().isEmpty()) return results; try { @@ -74,17 +50,17 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { 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()) { - ModelInfo match = uploadedMatch.get(); + String name = info.getName(); + + Optional listMatch = (modelListService != null) ? modelListService.findByFilename(name) : Optional.empty(); + if (listMatch.isPresent()) { + ModelInfo match = listMatch.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)); - info.setPopularity("๐Ÿง  AI VERIFIED OFFICIAL"); + info.setPopularity("๐Ÿ“‚ MODEL LIST MATCH"); } else { + // Fallback to AI prediction if not in list + if (aiService == null) aiService = new LocalAIService(); LocalAIService.Prediction prediction = aiService.predictProvider(info.getName()); info.setPopularity(prediction.getLabel()); @@ -147,13 +123,6 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { return null; } - private void applyContext(List res, String fn) { - if (fn == null) return; - String ctx = fn.toLowerCase().split("[_\\- ]")[0]; - if (ctx.length() < 3) return; - for (ModelInfo i : res) if (i.getName().toLowerCase().contains(ctx) && i.getPopularity().contains("Community")) i.setPopularity("๐Ÿ“Œ Context Match: " + ctx); - } - private void findModelsMetadata(Object obj, List res) { if (obj instanceof JSONObject) { JSONObject jo = (JSONObject) obj; @@ -212,16 +181,13 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { Matcher mdMatcher = ExpertPatterns.MD_LINK.matcher(v); while (mdMatcher.find()) { String name = mdMatcher.group(1); - String url = mdMatcher.group(2); + String url = fixUrl(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(); @@ -230,7 +196,21 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { } } } + if (type == null) type = ctx; + addModelInfo(res, new ModelInfo(type != null ? type : "checkpoints", name, url)); + } + } + // 2. Suche nach nackten URLs + Matcher nakedMatcher = ExpertPatterns.NAKED_URL.matcher(v); + while (nakedMatcher.find()) { + String url = fixUrl(nakedMatcher.group(1)); + String name = url; + if (name.contains("?")) name = name.substring(0, name.indexOf("?")); + if (name.contains("/")) name = name.substring(name.lastIndexOf("/") + 1); + + if (name.length() > 3) { + String type = hint != null ? inferTypeFromKey(hint) : null; if (type == null) type = ctx; addModelInfo(res, new ModelInfo(type != null ? type : "checkpoints", name, url)); } @@ -256,8 +236,9 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { String type = hint != null ? inferTypeFromKey(hint) : null; if (type == null) type = ctx; - // 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"; + // Verfeinerung basierend auf Dateinamen (Hรถchste Prioritรคt fรผr Spezialmodelle) + if (fLow.contains("qwen-image-lightning")) type = "unet"; + else 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"; @@ -316,14 +297,15 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { private void addModelInfo(List res, ModelInfo info) { for (ModelInfo e : res) { if (e.getName().equalsIgnoreCase(info.getName())) { - // 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")) { + boolean isEWeak = e.getType() == null || e.getType().equals("checkpoints") || e.getType().equals("clip"); + boolean isInfoStrong = info.getType() != null && (info.getType().equals("unet") || info.getType().equals("checkpoints")); + + if (isEWeak && isInfoStrong) { e.setType(info.getType()); e.setSave_path(info.getSave_path()); } @@ -333,4 +315,12 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { } res.add(info); } + + private String fixUrl(String url) { + if (url == null) return null; + if (url.contains("huggingface.co") && url.contains("/blob/")) { + return url.replace("/blob/", "/resolve/"); + } + return url; + } } 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 94fe5df..d08817e 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ConfigService.java @@ -84,7 +84,33 @@ public class ConfigService { 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 isBackgroundModeEnabled() { return settings.optBoolean("background_mode", false); } + public void setBackgroundModeEnabled(boolean enabled) { settings.put("background_mode", enabled); save(); } + + public boolean isShutdownAfterDownloadEnabled() { return settings.optBoolean("shutdown_after_download", false); } + public void setShutdownAfterDownloadEnabled(boolean enabled) { settings.put("shutdown_after_download", enabled); save(); } + + public void savePendingDownloads(String json) { + try { + Files.writeString(getFileInAppData("pending_downloads.json").toPath(), json, StandardCharsets.UTF_8); + } catch (Exception e) { + System.err.println("Error saving pending downloads: " + e.getMessage()); + } + } + + public String loadPendingDownloads() { + try { + File file = getFileInAppData("pending_downloads.json"); + if (file.exists()) { + return Files.readString(file.toPath(), StandardCharsets.UTF_8); + } + } catch (Exception e) { + System.err.println("Error loading pending downloads: " + e.getMessage()); + } + return null; + } + 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 d55ce57..9d93d19 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java +++ b/src/main/java/de/tki/comfymodels/service/impl/DefaultDownloadManager.java @@ -111,12 +111,24 @@ public class DefaultDownloadManager implements IDownloadManager { } private void downloadWithResume(ModelInfo info, Path targetFile, int index, BiConsumer statusUpdater) { + downloadWithResumeInternal(info, targetFile, index, statusUpdater, 0); + } + + private void downloadWithResumeInternal(ModelInfo info, Path targetFile, int index, BiConsumer statusUpdater, int retryCount) { try { if (waitForPauseAndCheckSelection(index, statusUpdater)) return; File file = targetFile.toFile(); String hfToken = configService.getHfToken(); + // Check disk space before starting + try { + long usableSpace = Files.getFileStore(targetFile.getParent().getRoot()).getUsableSpace(); + if (usableSpace < 10L * 1024 * 1024 * 1024) { // 10 GB Buffer + safeUpdateStatus(index, "โš ๏ธ Low Disk Space (<10GB)", statusUpdater); + } + } catch (Exception ignored) {} + if (file.exists() && file.length() < 10240) { file.delete(); } @@ -143,7 +155,25 @@ public class DefaultDownloadManager implements IDownloadManager { long totalRemoteSize = headResponse.headers().firstValueAsLong("Content-Length").orElse(0L); + // Double Check Disk Space against total size + try { + long usableSpace = Files.getFileStore(targetFile.getParent().getRoot()).getUsableSpace(); + if (totalRemoteSize > 0 && usableSpace < totalRemoteSize) { + safeUpdateStatus(index, "โŒ No Space (" + formatSize(usableSpace) + " < " + formatSize(totalRemoteSize) + ")", statusUpdater); + return; + } + } catch (Exception ignored) {} + if (existingFileSize > 0 && totalRemoteSize > 0 && existingFileSize == totalRemoteSize) { + // Verification of existing file + if (info.getName().endsWith(".safetensors") && (existingFileSize % 2 != 0)) { + if (retryCount < 1) { + safeUpdateStatus(index, "๐Ÿ”„ Fixing Corrupted File...", statusUpdater); + file.delete(); + downloadWithResumeInternal(info, targetFile, index, statusUpdater, retryCount + 1); + return; + } + } safeUpdateStatus(index, "โœ… Already exists", statusUpdater); return; } @@ -214,8 +244,18 @@ public class DefaultDownloadManager implements IDownloadManager { 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); + long finalSize = file.length(); + boolean sizeMismatch = totalBytes > 0 && finalSize < totalBytes; + boolean corruptedSafetensor = info.getName().endsWith(".safetensors") && (finalSize % 2 != 0); + + if ((sizeMismatch || corruptedSafetensor) && retryCount < 1) { + safeUpdateStatus(index, "๐Ÿ”„ Verification failed, redownloading...", statusUpdater); + file.delete(); + downloadWithResumeInternal(info, targetFile, index, statusUpdater, retryCount + 1); + } else if (sizeMismatch) { + safeUpdateStatus(index, "โŒ Incomplete (" + formatSize(finalSize) + "/" + formatSize(totalBytes) + ")", statusUpdater); + } else if (corruptedSafetensor) { + safeUpdateStatus(index, "โŒ Corrupted (Odd Size)", statusUpdater); } else { safeUpdateStatus(index, "โœ… Finished", statusUpdater); } 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 abb11bf..898c983 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ModelListService.java @@ -55,9 +55,29 @@ public class ModelListService { } public Optional findByFilename(String filename) { - return models.stream() - .filter(m -> filename.equalsIgnoreCase(m.getFilename())) + if (filename == null) return Optional.empty(); + String low = filename.toLowerCase(); + + // 1. Direct match + Optional direct = models.stream() + .filter(m -> filename.equalsIgnoreCase(m.getFilename()) || filename.equalsIgnoreCase(m.getName())) .findFirst(); + if (direct.isPresent()) return direct; + + // 2. Strip paths and extensions for fuzzy match + String lowStrip = low.contains("/") ? low.substring(low.lastIndexOf("/") + 1) : low; + if (lowStrip.contains("\\")) lowStrip = lowStrip.substring(lowStrip.lastIndexOf("\\") + 1); + String lowNoExt = lowStrip.contains(".") ? lowStrip.substring(0, lowStrip.lastIndexOf(".")) : lowStrip; + + final String finalLowStrip = lowStrip; + final String finalLowNoExt = lowNoExt; + + return models.stream().filter(m -> { + String mf = m.getFilename() != null ? m.getFilename().toLowerCase() : ""; + String mn = m.getName() != null ? m.getName().toLowerCase() : ""; + return mf.equals(finalLowStrip) || mn.equals(finalLowStrip) || + mf.startsWith(finalLowNoExt) || mn.startsWith(finalLowNoExt); + }).findFirst(); } private void loadFromStorage() { diff --git a/src/main/java/de/tki/comfymodels/service/impl/ModelValidator.java b/src/main/java/de/tki/comfymodels/service/impl/ModelValidator.java new file mode 100644 index 0000000..eaabb75 --- /dev/null +++ b/src/main/java/de/tki/comfymodels/service/impl/ModelValidator.java @@ -0,0 +1,142 @@ +package de.tki.comfymodels.service.impl; + +import de.tki.comfymodels.service.IModelValidator; +import org.json.JSONArray; +import org.json.JSONObject; +import org.springframework.stereotype.Service; + +import java.io.File; +import java.io.RandomAccessFile; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; + +@Service +public class ModelValidator implements IModelValidator { + + @Override + public ValidationResult validateFile(File file) { + String path = file.getAbsolutePath(); + + if (!file.exists()) return new ValidationResult(false, "File does not exist", path); + if (file.length() == 0) return new ValidationResult(false, "Empty file (0 bytes)", path); + + // Check for LFS Stubs (common when downloading from HF without proper LFS setup) + if (file.length() < 1024) { + try { + String content = Files.readString(file.toPath()); + if (content.contains("version https://git-lfs.github.com/spec/v1")) { + return new ValidationResult(false, "Hugging Face LFS Stub (Not the actual model)", path); + } + } catch (Exception ignored) {} + } + + String name = file.getName().toLowerCase(); + if (name.endsWith(".safetensors") || name.endsWith(".sft")) { + return validateSafetensors(file); + } + + // For other files, we just check if they are readable and have a sane size + if (file.length() < 1024 * 1024 && !name.endsWith(".yaml") && !name.endsWith(".json")) { + return new ValidationResult(false, "Suspiciously small for a model file", path); + } + + return new ValidationResult(true, "OK", path); + } + + private ValidationResult validateSafetensors(File file) { + try (RandomAccessFile raf = new RandomAccessFile(file, "r")) { + if (raf.length() < 8) return new ValidationResult(false, "File too small to have a header", file.getAbsolutePath()); + + byte[] headerLenBytes = new byte[8]; + raf.readFully(headerLenBytes); + + // Safetensors header length is a 64-bit little-endian unsigned integer + ByteBuffer buffer = ByteBuffer.wrap(headerLenBytes); + buffer.order(ByteOrder.LITTLE_ENDIAN); + long headerLen = buffer.getLong(); + + if (headerLen <= 0 || headerLen > raf.length() - 8) { + return new ValidationResult(false, "Invalid header length: " + headerLen, file.getAbsolutePath()); + } + + // Optional: Validate that the header is valid JSON + if (headerLen < 100 * 1024 * 1024) { // Don't try to read massive headers into memory + byte[] headerBytes = new byte[(int) headerLen]; + raf.readFully(headerBytes); + String headerJson = new String(headerBytes, StandardCharsets.UTF_8); + try { + JSONObject json = new JSONObject(headerJson); + long maxOffset = 0; + + for (String key : json.keySet()) { + if (key.equals("__metadata__")) continue; + + JSONObject tensorInfo = json.optJSONObject(key); + if (tensorInfo != null && tensorInfo.has("data_offsets")) { + JSONArray offsets = tensorInfo.getJSONArray("data_offsets"); + if (offsets.length() == 2) { + long start = offsets.getLong(0); + long end = offsets.getLong(1); + maxOffset = Math.max(maxOffset, end); + + // Check for alignment/element size issues + String dtype = tensorInfo.optString("dtype", ""); + if (!dtype.isEmpty()) { + int elementSize = getElementSize(dtype); + if (elementSize > 1) { + long length = end - start; + if (length % elementSize != 0) { + return new ValidationResult(false, + String.format("Tensor '%s' length (%d) is not a multiple of element size (%d) for dtype %s", + key, length, elementSize, dtype), file.getAbsolutePath()); + } + } + } + } + } + } + + long expectedSize = 8 + headerLen + maxOffset; + if (file.length() < expectedSize) { + return new ValidationResult(false, + String.format("File is truncated. Expected at least %d bytes but found %d", expectedSize, file.length()), + file.getAbsolutePath()); + } + + } catch (Exception e) { + return new ValidationResult(false, "Header is not valid JSON or has invalid structure: " + e.getMessage(), file.getAbsolutePath()); + } + } + + return new ValidationResult(true, "OK", file.getAbsolutePath()); + } catch (Exception e) { + return new ValidationResult(false, "Error reading file: " + e.getMessage(), file.getAbsolutePath()); + } + } + + private int getElementSize(String dtype) { + switch (dtype.toUpperCase()) { + case "F64": + case "I64": + case "U64": + return 8; + case "F32": + case "I32": + case "U32": + return 4; + case "F16": + case "BF16": + case "I16": + case "U16": + return 2; + case "I8": + case "U8": + case "BOOL": + return 1; + default: + return 1; // Unknown, assume 1 + } + } +} diff --git a/src/main/resources/known_good_urls.json b/src/main/resources/known_good_urls.json deleted file mode 100644 index 387ec75..0000000 --- a/src/main/resources/known_good_urls.json +++ /dev/null @@ -1,30 +0,0 @@ -{ - "flux1-schnell.sft": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors", - "flux1-dev.sft": "https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors", - "flux1-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors", - "ae.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors", - "ae.sft": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors", - "flux2-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.2-dev/resolve/main/flux2-dev.safetensors", - "clip_l.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors", - "t5xxl_fp16.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp16.safetensors", - "t5xxl_fp8_e4m3fn.safetensors": "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors", - "clip_g.safetensors": "https://huggingface.co/Comfy-Org/stable-diffusion-v3-medium/resolve/main/text_encoders/clip_g.safetensors", - "v1-5-pruned-emaonly.safetensors": "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors", - "sd_xl_base_1.0.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors", - "sd_xl_refiner_1.0.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/resolve/main/sd_xl_refiner_1.0.safetensors", - "sd3.5_large.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/sd3.5_large.safetensors", - "sd3.5_medium.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-3.5-medium/resolve/main/sd3.5_medium.safetensors", - "sd_xl_turbo_1.0_fp16.safetensors": "https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/sd_xl_turbo_1.0_fp16.safetensors", - "svd.safetensors": "https://huggingface.co/stabilityai/stable-video-diffusion-img2vid/resolve/main/svd.safetensors", - "svd_xt.safetensors": "https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt/resolve/main/svd_xt.safetensors", - "flux1-canny-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev/resolve/main/flux1-canny-dev.safetensors", - "flux1-depth-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev/resolve/main/flux1-depth-dev.safetensors", - "flux1-redux-dev.safetensors": "https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev/resolve/main/flux1-redux-dev.safetensors", - "hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors": "https://huggingface.co/Kwai-VGI/HunyuanVideo/resolve/main/hunyuan_video_720_cfgdistill_fp8_e4m3fn.safetensors", - "mochi_1_preview_fp8.safetensors": "https://huggingface.co/genmo/mochi-1-preview/resolve/main/mochi_1_preview_fp8.safetensors", - "ltx-video-2b-v0.9.safetensors": "https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltx-video-2b-v0.9.safetensors", - "vae-ft-mse-840000-ema-pruned.safetensors": "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors", - "sdxl_vae.safetensors": "https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl_vae.safetensors", - "RealESRGAN_x4plus.pth": "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4plus.pth", - "RealESRGAN_x4plus_anime_6B.pth": "https://huggingface.co/ai-forever/Real-ESRGAN/resolve/main/RealESRGAN_x4plus_anime_6B.pth" -} diff --git a/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java b/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java index 4750ea4..946fd1e 100644 --- a/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java +++ b/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java @@ -1,46 +1,49 @@ -package de.tki.comfymodels.service; +package de.tki.comfymodels; import de.tki.comfymodels.domain.ModelInfo; import de.tki.comfymodels.service.impl.ComfyModelAnalyzer; +import de.tki.comfymodels.service.impl.ModelListService; import org.junit.jupiter.api.Test; - +import java.lang.reflect.Field; import java.util.List; +import java.util.Optional; import static org.junit.jupiter.api.Assertions.*; public class ComfyModelAnalyzerTest { - private final ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer(); - @Test - public void testAnalyzeBasicWorkflow() { - String json = "{\"nodes\": [{\"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"flux1-schnell.sft\"]}]}"; + public void testAnalyzeBasicWorkflow() throws Exception { + ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer(); + + // Mock ModelListService + ModelListService mockListService = new ModelListService() { + @Override + public Optional findByFilename(String filename) { + if ("flux1-schnell.safetensors".equalsIgnoreCase(filename)) { + return Optional.of(new ModelInfo("checkpoints", "flux1-schnell.safetensors", + "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors")); + } + return Optional.empty(); + } + }; + + Field field = ComfyModelAnalyzer.class.getDeclaredField("modelListService"); + field.setAccessible(true); + field.set(analyzer, mockListService); + + String json = "{\"nodes\": [{\"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"flux1-schnell.safetensors\"]}]}"; List models = analyzer.analyze(json, "test.json"); - - assertNotNull(models); + assertFalse(models.isEmpty()); - - ModelInfo flux = models.stream() - .filter(m -> m.getName().equals("flux1-schnell.sft")) - .findFirst() - .orElse(null); - - assertNotNull(flux); - assertEquals("unet", flux.getType()); - assertEquals("https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors", flux.getUrl()); + assertEquals("flux1-schnell.safetensors", models.get(0).getName()); + assertEquals("https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors", models.get(0).getUrl()); } @Test - public void testAnalyzeLora() { - String json = "{\"nodes\": [{\"type\": \"LoraLoader\", \"widgets_values\": [\"my_cool_lora.safetensors\"]}]}"; - List models = analyzer.analyze(json, "test.json"); - - ModelInfo lora = models.stream() - .filter(m -> m.getName().equals("my_cool_lora.safetensors")) - .findFirst() - .orElse(null); - - assertNotNull(lora); - assertEquals("loras", lora.getType()); + public void testInvalidJson() { + ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer(); + List models = analyzer.analyze("invalid json", "test.json"); + assertTrue(models.isEmpty()); } } diff --git a/src/test/java/de/tki/comfymodels/service/ModelValidatorTest.java b/src/test/java/de/tki/comfymodels/service/ModelValidatorTest.java new file mode 100644 index 0000000..63dcc6b --- /dev/null +++ b/src/test/java/de/tki/comfymodels/service/ModelValidatorTest.java @@ -0,0 +1,136 @@ +package de.tki.comfymodels.service; + +import de.tki.comfymodels.service.impl.ModelValidator; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; + +import static org.junit.jupiter.api.Assertions.*; + +class ModelValidatorTest { + + private ModelValidator validator; + + @TempDir + Path tempDir; + + @BeforeEach + void setUp() { + validator = new ModelValidator(); + } + + @Test + void testNonExistentFile() { + File file = new File(tempDir.toFile(), "ghost.safetensors"); + IModelValidator.ValidationResult result = validator.validateFile(file); + assertFalse(result.ok); + assertTrue(result.message.contains("exist")); + } + + @Test + void testEmptyFile() throws IOException { + File file = new File(tempDir.toFile(), "empty.safetensors"); + file.createNewFile(); + IModelValidator.ValidationResult result = validator.validateFile(file); + assertFalse(result.ok); + assertTrue(result.message.contains("Empty")); + } + + @Test + void testLfsStub() throws IOException { + File file = new File(tempDir.toFile(), "stub.safetensors"); + String lfsContent = "version https://git-lfs.github.com/spec/v1\noid sha256:12345\nsize 1000"; + java.nio.file.Files.writeString(file.toPath(), lfsContent); + + IModelValidator.ValidationResult result = validator.validateFile(file); + assertFalse(result.ok); + assertTrue(result.message.contains("LFS Stub")); + } + + @Test + void testValidSafetensors() throws IOException { + File file = new File(tempDir.toFile(), "valid.safetensors"); + String headerJson = "{\"__metadata__\":{\"format\":\"pt\"}}"; + byte[] headerBytes = headerJson.getBytes(StandardCharsets.UTF_8); + + try (FileOutputStream fos = new FileOutputStream(file)) { + ByteBuffer buf = ByteBuffer.allocate(8); + buf.order(ByteOrder.LITTLE_ENDIAN); + buf.putLong(headerBytes.length); + fos.write(buf.array()); + fos.write(headerBytes); + fos.write(new byte[100]); // Fake data + } + + IModelValidator.ValidationResult result = validator.validateFile(file); + assertTrue(result.ok, "Should be valid: " + result.message); + } + + @Test + void testCorruptSafetensorsHeaderLength() throws IOException { + File file = new File(tempDir.toFile(), "corrupt_len.safetensors"); + + try (FileOutputStream fos = new FileOutputStream(file)) { + ByteBuffer buf = ByteBuffer.allocate(8); + buf.order(ByteOrder.LITTLE_ENDIAN); + buf.putLong(1000000); // Invalid length (much larger than file) + fos.write(buf.array()); + fos.write(new byte[10]); + } + + IModelValidator.ValidationResult result = validator.validateFile(file); + assertFalse(result.ok); + assertTrue(result.message.contains("Invalid header length")); + } + + @Test + void testTruncatedSafetensors() throws IOException { + File file = new File(tempDir.toFile(), "truncated.safetensors"); + // Header says we have 1000 bytes of data, but we only write 500 + String headerJson = "{\"tensor1\":{\"dtype\":\"F16\",\"shape\":[500],\"data_offsets\":[0,1000]}}"; + byte[] headerBytes = headerJson.getBytes(StandardCharsets.UTF_8); + + try (FileOutputStream fos = new FileOutputStream(file)) { + ByteBuffer buf = ByteBuffer.allocate(8); + buf.order(ByteOrder.LITTLE_ENDIAN); + buf.putLong(headerBytes.length); + fos.write(buf.array()); + fos.write(headerBytes); + fos.write(new byte[500]); // Only half the data + } + + IModelValidator.ValidationResult result = validator.validateFile(file); + // Currently this passes, but it should fail + assertFalse(result.ok, "Should detect truncation: " + result.message); + assertTrue(result.message.contains("truncated") || result.message.contains("size"), "Message should mention size or truncation: " + result.message); + } + + @Test + void testSafetensorsAlignmentError() throws IOException { + File file = new File(tempDir.toFile(), "alignment.safetensors"); + // Tensor says it's F16 (2 bytes), but length is 999 + String headerJson = "{\"tensor1\":{\"dtype\":\"F16\",\"shape\":[500],\"data_offsets\":[0,999]}}"; + byte[] headerBytes = headerJson.getBytes(StandardCharsets.UTF_8); + + try (FileOutputStream fos = new FileOutputStream(file)) { + ByteBuffer buf = ByteBuffer.allocate(8); + buf.order(ByteOrder.LITTLE_ENDIAN); + buf.putLong(headerBytes.length); + fos.write(buf.array()); + fos.write(headerBytes); + fos.write(new byte[1000]); // File has enough data, but header definition is misaligned + } + + IModelValidator.ValidationResult result = validator.validateFile(file); + assertFalse(result.ok, "Should detect alignment error"); + assertTrue(result.message.contains("multiple of element size"), "Message should mention element size: " + result.message); + } +}