227 lines
9.4 KiB
Java
227 lines
9.4 KiB
Java
package de.tki.comfymodels.service.impl;
|
|
|
|
import de.tki.comfymodels.domain.ModelInfo;
|
|
import org.springframework.beans.factory.annotation.Autowired;
|
|
import org.springframework.stereotype.Service;
|
|
|
|
import java.io.InputStream;
|
|
import java.io.OutputStream;
|
|
import java.io.IOException;
|
|
import java.nio.file.Files;
|
|
import java.nio.file.Path;
|
|
import java.nio.file.Paths;
|
|
import java.nio.file.StandardCopyOption;
|
|
import java.util.ArrayList;
|
|
import java.util.List;
|
|
import java.util.Map;
|
|
import java.util.TreeMap;
|
|
import java.util.function.LongConsumer;
|
|
import java.util.stream.Stream;
|
|
|
|
@Service
|
|
public class ArchiveService {
|
|
|
|
private final ConfigService configService;
|
|
private final PathResolver pathResolver;
|
|
|
|
@Autowired
|
|
public ArchiveService(ConfigService configService, PathResolver pathResolver) {
|
|
this.configService = configService;
|
|
this.pathResolver = pathResolver;
|
|
}
|
|
|
|
public String normalizeFolder(String folder) {
|
|
return pathResolver.stripRedundantPrefixes(folder);
|
|
}
|
|
|
|
public Map<String, List<ModelInfo>> getModelsGroupedByFolder() {
|
|
return scanDirectory(configService.getModelsPath(), true);
|
|
}
|
|
|
|
public Map<String, List<ModelInfo>> getArchivedModelsGroupedByFolder() {
|
|
return scanDirectory(configService.getArchivePath(), false);
|
|
}
|
|
|
|
private Map<String, List<ModelInfo>> scanDirectory(String pathStr, boolean excludeArchived) {
|
|
Map<String, List<ModelInfo>> grouped = new TreeMap<>();
|
|
if (pathStr == null || pathStr.isEmpty()) return grouped;
|
|
|
|
Path root = Paths.get(pathStr);
|
|
if (!Files.exists(root)) {
|
|
System.err.println("Scan directory does not exist: " + pathStr);
|
|
return grouped;
|
|
}
|
|
|
|
try (Stream<Path> walk = Files.walk(root)) {
|
|
walk.filter(p -> {
|
|
try {
|
|
// Ignore hidden directories and .venv / venv
|
|
Path relativeToRoot = root.relativize(p);
|
|
for (Path part : relativeToRoot) {
|
|
String name = part.toString();
|
|
if (name.startsWith(".") || name.equalsIgnoreCase("venv") || name.equalsIgnoreCase(".venv") || name.equalsIgnoreCase("__pycache__")) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
} catch (IllegalArgumentException e) {
|
|
return false;
|
|
}
|
|
})
|
|
.filter(p -> Files.isRegularFile(p))
|
|
.filter(this::isSupportedModel)
|
|
.filter(p -> !excludeArchived || !isAlreadyArchived(p))
|
|
.forEach(p -> {
|
|
if (!Files.exists(p)) return; // Double check existence
|
|
|
|
Path relative = root.relativize(p);
|
|
String folder = relative.getParent() != null ? relative.getParent().toString().replace("\\", "/") : "root";
|
|
|
|
ModelInfo info = new ModelInfo();
|
|
info.setName(p.getFileName().toString());
|
|
info.setSave_path(folder);
|
|
info.setType(folder);
|
|
|
|
try {
|
|
long bytes = Files.size(p);
|
|
double mb = bytes / (1024.0 * 1024.0);
|
|
if (mb > 1024) {
|
|
info.setSize(String.format("%.2f GB", mb / 1024.0));
|
|
} else {
|
|
info.setSize(String.format("%.2f MB", mb));
|
|
}
|
|
} catch (IOException e) {
|
|
info.setSize("Unknown");
|
|
}
|
|
|
|
grouped.computeIfAbsent(folder, k -> new ArrayList<>()).add(info);
|
|
});
|
|
} catch (IOException e) {
|
|
System.err.println("Error walking directory " + pathStr + ": " + e.getMessage());
|
|
e.printStackTrace();
|
|
}
|
|
|
|
return grouped;
|
|
}
|
|
|
|
private boolean isSupportedModel(Path 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");
|
|
}
|
|
|
|
private boolean isAlreadyArchived(Path p) {
|
|
String archivePathStr = configService.getArchivePath();
|
|
if (archivePathStr == null || archivePathStr.isEmpty()) return false;
|
|
Path archivePath = Paths.get(archivePathStr).toAbsolutePath();
|
|
return p.toAbsolutePath().startsWith(archivePath);
|
|
}
|
|
|
|
public void moveToArchiveWithProgress(String relativePath, LongConsumer progressUpdate) throws IOException {
|
|
String modelsPath = configService.getModelsPath();
|
|
String archivePath = configService.getArchivePath();
|
|
|
|
if (modelsPath == null || archivePath == null) throw new IOException("Paths not configured");
|
|
|
|
Path source = Paths.get(modelsPath, relativePath);
|
|
Path target = Paths.get(archivePath, relativePath);
|
|
|
|
if (!Files.exists(source)) throw new IOException("Source not found: " + source);
|
|
|
|
Files.createDirectories(target.getParent());
|
|
moveWithProgress(source, target, progressUpdate);
|
|
}
|
|
|
|
public void moveToArchive(String relativePath) throws IOException {
|
|
moveToArchiveWithProgress(relativePath, null);
|
|
}
|
|
|
|
public void moveToArchiveWithProgress(String folder, String filename, LongConsumer progressUpdate) throws IOException {
|
|
String normFolder = normalizeFolder(folder);
|
|
String relPath = "root".equals(normFolder) ? filename : Paths.get(normFolder, filename).toString();
|
|
moveToArchiveWithProgress(relPath, progressUpdate);
|
|
}
|
|
|
|
public void moveToArchive(String folder, String filename) throws IOException {
|
|
moveToArchiveWithProgress(folder, filename, null);
|
|
}
|
|
|
|
public boolean restoreFromArchiveWithProgress(String folder, String filename, LongConsumer progressUpdate) {
|
|
String modelsPathStr = configService.getModelsPath();
|
|
String archivePathStr = configService.getArchivePath();
|
|
|
|
if (modelsPathStr == null || archivePathStr == null) return false;
|
|
|
|
Path modelsPath = Paths.get(modelsPathStr);
|
|
Path archivePath = Paths.get(archivePathStr);
|
|
|
|
String normFolder = normalizeFolder(folder);
|
|
|
|
// Primary location with case-insensitive resolution
|
|
Path archivedFolder = "root".equals(normFolder) ? archivePath : pathResolver.resolveCaseInsensitiveRecursive(archivePath, normFolder);
|
|
Path archived = archivedFolder.resolve(filename);
|
|
|
|
Path targetFolder = "root".equals(normFolder) ? modelsPath : modelsPath.resolve(normFolder);
|
|
Path target = targetFolder.resolve(filename);
|
|
|
|
// Fallback: If not at primary, check root of archive
|
|
if (!Files.exists(archived)) {
|
|
Path rootArchived = archivePath.resolve(filename);
|
|
if (Files.exists(rootArchived)) {
|
|
archived = rootArchived;
|
|
} else {
|
|
// Secondary Fallback: Check if it exists ANYWHERE in the archive (shallow search)
|
|
try (Stream<Path> walk = Files.walk(archivePath, 2)) {
|
|
java.util.Optional<Path> found = walk.filter(Files::isRegularFile)
|
|
.filter(p -> p.getFileName().toString().equalsIgnoreCase(filename))
|
|
.findFirst();
|
|
if (found.isPresent()) {
|
|
archived = found.get();
|
|
} else {
|
|
System.err.println("Restore failed: " + filename + " not found in archive (checked " + normFolder + " and root)");
|
|
return false;
|
|
}
|
|
} catch (IOException e) {
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
|
|
try {
|
|
if (Files.exists(archived) && Files.size(archived) > 0) {
|
|
Files.createDirectories(target.getParent());
|
|
moveWithProgress(archived, target, progressUpdate);
|
|
return true;
|
|
}
|
|
} catch (IOException e) {
|
|
System.err.println("Error moving file from archive: " + e.getMessage());
|
|
e.printStackTrace();
|
|
}
|
|
return false;
|
|
}
|
|
|
|
public boolean restoreFromArchive(String folder, String filename) {
|
|
return restoreFromArchiveWithProgress(folder, filename, null);
|
|
}
|
|
|
|
private void moveWithProgress(Path source, Path target, LongConsumer progressUpdate) throws IOException {
|
|
// Optimization: If on same file store, move is instant
|
|
if (Files.getFileStore(source).equals(Files.getFileStore(target.getParent()))) {
|
|
Files.move(source, target, StandardCopyOption.REPLACE_EXISTING);
|
|
if (progressUpdate != null) progressUpdate.accept(Files.size(target));
|
|
return;
|
|
}
|
|
|
|
// Cross-drive move: Copy with progress, then delete
|
|
try (InputStream in = Files.newInputStream(source);
|
|
OutputStream out = Files.newOutputStream(target)) {
|
|
byte[] buffer = new byte[1024 * 64]; // 64KB buffer
|
|
int bytesRead;
|
|
while ((bytesRead = in.read(buffer)) != -1) {
|
|
out.write(buffer, 0, bytesRead);
|
|
if (progressUpdate != null) progressUpdate.accept((long) bytesRead);
|
|
}
|
|
}
|
|
Files.delete(source);
|
|
}
|
|
}
|