Thomas
5c6f37741f
- Block 1: Refactored path logic into a centralized PathResolver for OS-agnostic and relative path handling. - Block 2: Improved platform compatibility (shutdown, macOS menu, System Tray fallback). - Block 3: Added Power-User features (Model Updater, Smart Cleanup, Thumbnails, Node Sync, Remote Status). - Fixes: Recursive subfolder scanning, archive detection logic, and UI progress throttling.
392 lines
19 KiB
Java
392 lines
19 KiB
Java
package de.tki.comfymodels.service.impl;
|
|
|
|
import de.tki.comfymodels.domain.ModelInfo;
|
|
import de.tki.comfymodels.service.IModelSearchService;
|
|
import org.json.JSONArray;
|
|
import org.json.JSONObject;
|
|
import org.springframework.beans.factory.annotation.Autowired;
|
|
import org.springframework.stereotype.Service;
|
|
|
|
import java.net.URI;
|
|
import java.net.URLEncoder;
|
|
import java.net.http.HttpClient;
|
|
import java.net.http.HttpRequest;
|
|
import java.net.http.HttpResponse;
|
|
import java.nio.charset.StandardCharsets;
|
|
import java.util.*;
|
|
import java.util.concurrent.ExecutorService;
|
|
import java.util.concurrent.Executors;
|
|
import java.util.concurrent.atomic.AtomicInteger;
|
|
import java.util.function.BiConsumer;
|
|
|
|
@Service
|
|
public class ModelSearchService implements IModelSearchService {
|
|
|
|
@Autowired
|
|
private ConfigService configService;
|
|
|
|
@Autowired
|
|
private GeminiAIService geminiService;
|
|
|
|
@Autowired
|
|
private ModelListService modelListService;
|
|
|
|
private static final Set<String> OFFICIAL_AUTHORS = new HashSet<>(Arrays.asList(
|
|
"black-forest-labs", "stabilityai", "runwayml", "comfyanonymous", "Comfy-Org",
|
|
"lllyasviel", "Kwai-Kolors", "Kwai-VGI", "InstantX", "ByteDance", "apple",
|
|
"google", "facebook", "TencentARC", "DeepFloyd", "microsoft", "nvidia", "fal", "genmo",
|
|
"ostris", "XLabs-AI", "THUDM", "city96", "Kijai", "BartoszGawlik"
|
|
));
|
|
|
|
private final HttpClient httpClient = HttpClient.newBuilder()
|
|
.followRedirects(HttpClient.Redirect.ALWAYS)
|
|
.build();
|
|
|
|
private final ExecutorService searchExecutor = Executors.newFixedThreadPool(4);
|
|
|
|
@Override
|
|
public void searchOnline(List<ModelInfo> modelsToDownload, boolean[] selectedIndices, String workflowContext, String fileName,
|
|
BiConsumer<Integer, String> onStatusUpdate,
|
|
BiConsumer<Integer, ModelInfo> onModelFound,
|
|
Runnable onFinished) {
|
|
if (modelsToDownload == null) {
|
|
if (onFinished != null) onFinished.run();
|
|
return;
|
|
}
|
|
|
|
List<Integer> targetIndices = new ArrayList<>();
|
|
for (int i = 0; i < modelsToDownload.size(); i++) {
|
|
if (selectedIndices != null && i < selectedIndices.length && !selectedIndices[i]) continue;
|
|
targetIndices.add(i);
|
|
}
|
|
|
|
if (targetIndices.isEmpty()) {
|
|
if (onFinished != null) onFinished.run();
|
|
return;
|
|
}
|
|
|
|
AtomicInteger remaining = new AtomicInteger(targetIndices.size());
|
|
|
|
for (int index : targetIndices) {
|
|
ModelInfo info = modelsToDownload.get(index);
|
|
searchExecutor.submit(() -> {
|
|
try {
|
|
performSearch(info, index, fileName, workflowContext, onStatusUpdate, onModelFound);
|
|
} catch (Exception e) {
|
|
onStatusUpdate.accept(index, "Error: " + e.getMessage());
|
|
} finally {
|
|
if (remaining.decrementAndGet() == 0 && onFinished != null) {
|
|
onFinished.run();
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
@Autowired
|
|
private de.tki.comfymodels.service.impl.ModelHashRegistry hashRegistry;
|
|
|
|
@Autowired
|
|
private de.tki.comfymodels.service.IModelValidator modelValidator;
|
|
|
|
protected String getHfApiBaseUrl() {
|
|
return "https://huggingface.co/api";
|
|
}
|
|
|
|
protected String getCivitaiApiBaseUrl() {
|
|
return "https://civitai.com/api/v1";
|
|
}
|
|
|
|
protected String getHfResolveBaseUrl() {
|
|
return "https://huggingface.co";
|
|
}
|
|
|
|
private void performSearch(ModelInfo info, int index, String fileName, String workflowContext,
|
|
BiConsumer<Integer, String> onStatusUpdate,
|
|
BiConsumer<Integer, ModelInfo> onModelFound) {
|
|
// Priority 1: User defined Model List
|
|
Optional<ModelInfo> manualMatch = modelListService.findByFilename(info.getName());
|
|
if (manualMatch.isPresent() && !manualMatch.get().getUrl().equals("MISSING")) {
|
|
onStatusUpdate.accept(index, "📂 Found in Model List");
|
|
if (validateAndSetUrl(info, index, manualMatch.get().getUrl(), "📂 USER DEFINED", onStatusUpdate, onModelFound)) return;
|
|
}
|
|
|
|
// Priority 2: Civitai Hash Lookup
|
|
String modelsPath = configService.getModelsPath();
|
|
java.io.File localFile = new java.io.File(modelsPath, (info.getSave_path() != null ? info.getSave_path() : (info.getType() != null ? info.getType() : de.tki.comfymodels.domain.ModelFolder.CHECKPOINTS.getDefaultFolderName())) + java.io.File.separator + info.getName());
|
|
if (localFile.exists()) {
|
|
onStatusUpdate.accept(index, "🔍 Hashing for Civitai...");
|
|
String hash = hashRegistry.getOrCalculateHash(localFile);
|
|
if (hash != null) {
|
|
if (fetchCivitaiUrlByHash(info, index, hash, onStatusUpdate, onModelFound)) return;
|
|
}
|
|
}
|
|
|
|
onStatusUpdate.accept(index, "✨ Gemini Scouting...");
|
|
String aiHint = geminiService.discoverBestRepo(info.getName(), fileName, workflowContext);
|
|
if (aiHint != null && !aiHint.equalsIgnoreCase("UNKNOWN")) {
|
|
if (aiHint.startsWith("http")) {
|
|
if (validateAndSetUrl(info, index, aiHint, "✨ AI DIRECT", onStatusUpdate, onModelFound)) return;
|
|
}
|
|
onStatusUpdate.accept(index, "🔍 Validating Repo: " + aiHint);
|
|
if (fetchHuggingFaceUrlInSpecificRepo(info, index, aiHint, onStatusUpdate, onModelFound)) return;
|
|
}
|
|
|
|
String modelName = info.getName();
|
|
String cleanName = modelName.replaceAll("(_fp8|_fp16|_bf16|_v\\d+|\\d+v|_fix|\\.safetensors|\\.sft|\\.ckpt)", "");
|
|
|
|
LinkedHashSet<String> attempts = new LinkedHashSet<>();
|
|
attempts.add(modelName);
|
|
attempts.add(cleanName);
|
|
|
|
// Official HF
|
|
for (String query : attempts) {
|
|
if (query.length() < 3) continue;
|
|
onStatusUpdate.accept(index, "🔍 Searching Official: " + query);
|
|
if (fetchHuggingFaceUrl(info, index, query, true, onStatusUpdate, onModelFound)) return;
|
|
}
|
|
|
|
// Community HF
|
|
for (String query : attempts) {
|
|
if (query.length() < 3) continue;
|
|
onStatusUpdate.accept(index, "🔍 Searching Community: " + query);
|
|
if (fetchHuggingFaceUrl(info, index, query, false, onStatusUpdate, onModelFound)) return;
|
|
}
|
|
|
|
// Civitai Fallback
|
|
for (String query : attempts) {
|
|
if (query.length() < 3) continue;
|
|
onStatusUpdate.accept(index, "🔍 Searching Civitai: " + query);
|
|
if (fetchCivitaiUrl(info, index, query, onStatusUpdate, onModelFound)) return;
|
|
}
|
|
|
|
onStatusUpdate.accept(index, "❌ No trusted match found");
|
|
}
|
|
|
|
@Override
|
|
public Optional<ModelInfo> checkForUpdate(ModelInfo current) {
|
|
String modelsPath = configService.getModelsPath();
|
|
String subPath = (current.getSave_path() != null ? current.getSave_path() : (current.getType() != null ? current.getType() : de.tki.comfymodels.domain.ModelFolder.CHECKPOINTS.getDefaultFolderName()));
|
|
java.io.File localFile = new java.io.File(modelsPath, subPath + java.io.File.separator + current.getName());
|
|
|
|
if (!localFile.exists()) {
|
|
return Optional.empty();
|
|
}
|
|
|
|
String hash = hashRegistry.getOrCalculateHash(localFile);
|
|
if (hash == null) return Optional.empty();
|
|
|
|
try {
|
|
String hashUrl = getCivitaiApiBaseUrl() + "/model-versions/by-hash/" + hash;
|
|
HttpResponse<String> hashRes = httpClient.send(HttpRequest.newBuilder().uri(URI.create(hashUrl)).build(), HttpResponse.BodyHandlers.ofString());
|
|
if (hashRes.statusCode() != 200) return Optional.empty();
|
|
|
|
JSONObject currentVersion = new JSONObject(hashRes.body());
|
|
int modelId = currentVersion.optInt("modelId", -1);
|
|
int currentVersionId = currentVersion.optInt("id", -1);
|
|
if (modelId == -1 || currentVersionId == -1) return Optional.empty();
|
|
|
|
String modelUrl = getCivitaiApiBaseUrl() + "/models/" + modelId;
|
|
HttpResponse<String> modelRes = httpClient.send(HttpRequest.newBuilder().uri(URI.create(modelUrl)).build(), HttpResponse.BodyHandlers.ofString());
|
|
if (modelRes.statusCode() != 200) return Optional.empty();
|
|
|
|
JSONObject modelData = new JSONObject(modelRes.body());
|
|
JSONArray versions = modelData.optJSONArray("modelVersions");
|
|
if (versions == null || versions.length() == 0) return Optional.empty();
|
|
|
|
JSONObject latestVersion = versions.getJSONObject(0);
|
|
int latestVersionId = latestVersion.optInt("id", -1);
|
|
|
|
if (latestVersionId != -1 && latestVersionId != currentVersionId) {
|
|
JSONArray files = latestVersion.optJSONArray("files");
|
|
if (files != null) {
|
|
for (int i = 0; i < files.length(); i++) {
|
|
JSONObject file = files.getJSONObject(i);
|
|
if (file.optBoolean("primary", false) || files.length() == 1) {
|
|
ModelInfo updated = new ModelInfo();
|
|
updated.setType(current.getType());
|
|
updated.setSave_path(current.getSave_path());
|
|
updated.setName(file.getString("name"));
|
|
updated.setUrl(file.getString("downloadUrl"));
|
|
|
|
long size = getRemoteSize(updated.getUrl());
|
|
updated.setByteSize(size);
|
|
updated.setSize(formatSize(size));
|
|
updated.setPopularity("⭐ UPDATE: " + latestVersion.optString("name", "New Version"));
|
|
return Optional.of(updated);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} catch (Exception e) {
|
|
// Fail silently
|
|
}
|
|
return Optional.empty();
|
|
}
|
|
|
|
private boolean validateAndSetUrl(ModelInfo info, int rowIndex, String url, String popPrefix,
|
|
BiConsumer<Integer, String> onStatusUpdate,
|
|
BiConsumer<Integer, ModelInfo> onModelFound) {
|
|
if (url == null) return false;
|
|
|
|
if (url.contains("huggingface.co") && url.contains("/blob/")) {
|
|
url = url.replace("/blob/", "/resolve/");
|
|
}
|
|
|
|
long size = getRemoteSize(url);
|
|
if (size > 100) {
|
|
info.setUrl(url);
|
|
info.setPopularity(popPrefix);
|
|
info.setByteSize(size);
|
|
info.setSize(formatSize(size));
|
|
onModelFound.accept(rowIndex, info);
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
@Override
|
|
public long getRemoteSize(String url) {
|
|
try {
|
|
if (url.contains("huggingface.co") && url.contains("/resolve/")) {
|
|
String apiUrl = url.replace("/resolve/", "/api/models/").replace("/main/", "/file/");
|
|
if (!apiUrl.contains("/api/models/")) {
|
|
// Fallback for direct resolve URLs
|
|
apiUrl = getHfApiBaseUrl() + "/models/" + url.split("huggingface.co/")[1].replace("/resolve/main/", "/file/");
|
|
}
|
|
HttpRequest.Builder apiBuilder = HttpRequest.newBuilder().uri(URI.create(apiUrl)).GET().header("User-Agent", "Mozilla/5.0");
|
|
String token = configService.getHfToken();
|
|
if (!token.isEmpty()) apiBuilder.header("Authorization", "Bearer " + token);
|
|
|
|
HttpResponse<String> apiRes = httpClient.send(apiBuilder.build(), HttpResponse.BodyHandlers.ofString());
|
|
if (apiRes.statusCode() == 200) {
|
|
JSONObject fileData = new JSONObject(apiRes.body());
|
|
if (fileData.has("size")) return fileData.getLong("size");
|
|
}
|
|
}
|
|
|
|
HttpRequest.Builder builder = HttpRequest.newBuilder().uri(URI.create(url)).method("HEAD", HttpRequest.BodyPublishers.noBody()).header("User-Agent", "Mozilla/5.0");
|
|
String token = configService.getHfToken();
|
|
if (url.contains("huggingface.co") && !token.isEmpty()) builder.header("Authorization", "Bearer " + token);
|
|
|
|
HttpResponse<Void> res = httpClient.send(builder.build(), HttpResponse.BodyHandlers.discarding());
|
|
if (res.statusCode() == 401 || res.statusCode() == 403) return -401;
|
|
|
|
return res.headers().firstValueAsLong("Content-Length").orElse(-1L);
|
|
} catch (Exception e) {
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public String formatSize(long bytes) {
|
|
if (bytes == -401) return "🔒 Auth Required";
|
|
if (bytes <= 0) return "Unknown";
|
|
if (bytes >= 1024L * 1024L * 1024L) return String.format(Locale.US, "%.2f GB", bytes / 1024.0 / 1024.0 / 1024.0);
|
|
return String.format(Locale.US, "%.2f MB", bytes / 1024.0 / 1024.0);
|
|
}
|
|
|
|
private boolean fetchHuggingFaceUrlInSpecificRepo(ModelInfo info, int rowIndex, String repoId,
|
|
BiConsumer<Integer, String> onStatusUpdate,
|
|
BiConsumer<Integer, ModelInfo> onModelFound) {
|
|
try {
|
|
String treeUrl = getHfApiBaseUrl() + "/models/" + repoId + "/tree/main?recursive=true";
|
|
HttpRequest.Builder builder = HttpRequest.newBuilder().uri(URI.create(treeUrl)).GET();
|
|
String token = configService.getHfToken();
|
|
if (!token.isEmpty()) builder.header("Authorization", "Bearer " + token);
|
|
|
|
HttpResponse<String> response = httpClient.send(builder.build(), HttpResponse.BodyHandlers.ofString());
|
|
if (response.statusCode() == 200) {
|
|
JSONArray files = new JSONArray(response.body());
|
|
String targetName = info.getName().toLowerCase();
|
|
|
|
for (int j = 0; j < files.length(); j++) {
|
|
JSONObject file = files.getJSONObject(j);
|
|
String path = file.optString("path");
|
|
if ("file".equals(file.optString("type"))) {
|
|
String fileNameOnly = path.contains("/") ? path.substring(path.lastIndexOf("/") + 1) : path;
|
|
if (path.equalsIgnoreCase(info.getName()) || fileNameOnly.equalsIgnoreCase(targetName)) {
|
|
String downloadUrl = getHfResolveBaseUrl() + "/" + repoId + "/resolve/main/" + path;
|
|
return validateAndSetUrl(info, rowIndex, downloadUrl, "✨ AI RECURSIVE", onStatusUpdate, onModelFound);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} catch (Exception e) {}
|
|
return false;
|
|
}
|
|
|
|
private boolean fetchHuggingFaceUrl(ModelInfo info, int rowIndex, String query, boolean officialOnly,
|
|
BiConsumer<Integer, String> onStatusUpdate,
|
|
BiConsumer<Integer, ModelInfo> onModelFound) {
|
|
try {
|
|
String searchUrl = getHfApiBaseUrl() + "/models?search=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=downloads&direction=-1&limit=100";
|
|
HttpResponse<String> response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString());
|
|
if (response.statusCode() == 200) {
|
|
JSONArray models = new JSONArray(response.body());
|
|
for (int i = 0; i < models.length(); i++) {
|
|
String modelId = models.getJSONObject(i).getString("id");
|
|
String author = modelId.split("/")[0];
|
|
if (officialOnly && !OFFICIAL_AUTHORS.contains(author)) continue;
|
|
if (fetchHuggingFaceUrlInSpecificRepo(info, rowIndex, modelId, onStatusUpdate, onModelFound)) return true;
|
|
}
|
|
}
|
|
} catch (Exception e) {}
|
|
return false;
|
|
}
|
|
|
|
private boolean fetchCivitaiUrlByHash(ModelInfo info, int rowIndex, String hash,
|
|
BiConsumer<Integer, String> onStatusUpdate,
|
|
BiConsumer<Integer, ModelInfo> onModelFound) {
|
|
try {
|
|
String searchUrl = getCivitaiApiBaseUrl() + "/model-versions/by-hash/" + hash;
|
|
HttpResponse<String> response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString());
|
|
if (response.statusCode() == 200) {
|
|
JSONObject version = new JSONObject(response.body());
|
|
JSONArray files = version.optJSONArray("files");
|
|
if (files != null) {
|
|
for (int j = 0; j < files.length(); j++) {
|
|
JSONObject file = files.getJSONObject(j);
|
|
if (file.optString("name").equalsIgnoreCase(info.getName()) || file.optBoolean("primary", false)) {
|
|
return validateAndSetUrl(info, rowIndex, file.optString("downloadUrl"), "🛡️ Civitai (Hash Match)", onStatusUpdate, onModelFound);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} catch (Exception e) {}
|
|
return false;
|
|
}
|
|
|
|
private boolean fetchCivitaiUrl(ModelInfo info, int rowIndex, String query,
|
|
BiConsumer<Integer, String> onStatusUpdate,
|
|
BiConsumer<Integer, ModelInfo> onModelFound) {
|
|
try {
|
|
String searchUrl = getCivitaiApiBaseUrl() + "/models?query=" + URLEncoder.encode(query, StandardCharsets.UTF_8) + "&sort=Most+Downloaded&limit=20";
|
|
HttpResponse<String> response = httpClient.send(HttpRequest.newBuilder().uri(URI.create(searchUrl)).build(), HttpResponse.BodyHandlers.ofString());
|
|
if (response.statusCode() == 200) {
|
|
JSONObject json = new JSONObject(response.body());
|
|
JSONArray items = json.optJSONArray("items");
|
|
if (items != null) {
|
|
for (int i = 0; i < items.length(); i++) {
|
|
JSONArray versions = items.getJSONObject(i).optJSONArray("modelVersions");
|
|
if (versions != null) {
|
|
for (int v = 0; v < versions.length(); v++) {
|
|
JSONArray files = versions.getJSONObject(v).optJSONArray("files");
|
|
if (files != null) {
|
|
for (int j = 0; j < files.length(); j++) {
|
|
JSONObject file = files.getJSONObject(j);
|
|
if (file.optString("name").equalsIgnoreCase(info.getName()) || file.optString("name").toLowerCase().contains(query.toLowerCase())) {
|
|
return validateAndSetUrl(info, rowIndex, file.optString("downloadUrl"), "🛡️ Civitai", onStatusUpdate, onModelFound);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} catch (Exception e) {}
|
|
return false;
|
|
}
|
|
}
|