Add CivitAI hash lookup with progressive UI feedback
- Added lookup_civitai_by_hash() to query CivitAI API by file hash - Modified lookup-sources endpoint to process one model at a time - UI now shows real-time progress with found/not-found counts - Added Stop button to cancel ongoing lookup - Computes SHA256 hashes for models that don't have them This allows identifying models that were downloaded outside of tracked sources and finding their CivitAI pages automatically. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Ben Finklea
co-authored by
Claude Opus 4.5
Claude Opus 4.5
parent
77e81efa2d
commit
261d3c598c
@@ -177,6 +177,8 @@ class ModelSyncPanel {
|
||||
<option value="auto">Auto-Update Enabled</option>
|
||||
</select>
|
||||
<button class="modelsync-btn" data-action="scan">Scan</button>
|
||||
<button class="modelsync-btn primary" data-action="lookup">Lookup Sources</button>
|
||||
<button class="modelsync-btn" data-action="stop-lookup" style="display:none;background:#c00;">Stop</button>
|
||||
<button class="modelsync-btn" data-action="check">Check Updates</button>
|
||||
</div>
|
||||
<div class="modelsync-models-list"></div>
|
||||
@@ -278,6 +280,12 @@ class ModelSyncPanel {
|
||||
case "scan":
|
||||
await this.scanModels();
|
||||
break;
|
||||
case "lookup":
|
||||
await this.lookupSources();
|
||||
break;
|
||||
case "stop-lookup":
|
||||
this.stopLookup();
|
||||
break;
|
||||
case "check":
|
||||
await this.checkUpdates();
|
||||
break;
|
||||
@@ -305,6 +313,83 @@ class ModelSyncPanel {
|
||||
}
|
||||
}
|
||||
|
||||
async lookupSources() {
|
||||
this.lookupRunning = true;
|
||||
let foundCount = 0;
|
||||
let notFoundCount = 0;
|
||||
let processedCount = 0;
|
||||
|
||||
// Show stop button, hide lookup button
|
||||
const lookupBtn = this.element?.querySelector('[data-action="lookup"]');
|
||||
const stopBtn = this.element?.querySelector('[data-action="stop-lookup"]');
|
||||
if (lookupBtn) lookupBtn.style.display = 'none';
|
||||
if (stopBtn) stopBtn.style.display = 'inline-block';
|
||||
|
||||
const processNext = async () => {
|
||||
if (!this.lookupRunning) {
|
||||
if (lookupBtn) lookupBtn.style.display = 'inline-block';
|
||||
if (stopBtn) stopBtn.style.display = 'none';
|
||||
this.setStatus(`Lookup stopped. Found ${foundCount}, not found ${notFoundCount}.`);
|
||||
await this.refreshData();
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const res = await fetch("/modelsync/lookup-sources", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({}),
|
||||
});
|
||||
const data = await res.json();
|
||||
|
||||
if (data.status === "complete") {
|
||||
this.lookupRunning = false;
|
||||
if (lookupBtn) lookupBtn.style.display = 'inline-block';
|
||||
if (stopBtn) stopBtn.style.display = 'none';
|
||||
this.setStatus(`Lookup complete. Found ${foundCount} on CivitAI, ${notFoundCount} not found.`);
|
||||
await this.refreshData();
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.status === "found") {
|
||||
foundCount++;
|
||||
processedCount++;
|
||||
} else if (data.status === "not_found") {
|
||||
notFoundCount++;
|
||||
processedCount++;
|
||||
} else if (data.status === "hashing") {
|
||||
// Just hashed, will lookup next call
|
||||
}
|
||||
|
||||
const progress = data.total > 0 ? Math.round(((data.total - data.remaining) / data.total) * 100) : 0;
|
||||
this.setStatus(`[${progress}%] ${data.message} (${foundCount} found, ${data.remaining} remaining)`);
|
||||
|
||||
// Update UI periodically
|
||||
if (processedCount % 5 === 0) {
|
||||
await this.refreshData();
|
||||
}
|
||||
|
||||
// Continue to next model
|
||||
setTimeout(processNext, 100);
|
||||
|
||||
} catch (error) {
|
||||
this.lookupRunning = false;
|
||||
if (lookupBtn) lookupBtn.style.display = 'inline-block';
|
||||
if (stopBtn) stopBtn.style.display = 'none';
|
||||
this.setStatus("Lookup error: " + error.message);
|
||||
await this.refreshData();
|
||||
}
|
||||
};
|
||||
|
||||
this.setStatus("Starting source lookup...");
|
||||
processNext();
|
||||
}
|
||||
|
||||
stopLookup() {
|
||||
this.lookupRunning = false;
|
||||
this.setStatus("Stopping lookup...");
|
||||
}
|
||||
|
||||
async checkUpdates() {
|
||||
this.setStatus("Checking for updates...");
|
||||
try {
|
||||
|
||||
@@ -342,3 +342,98 @@ class VersionTracker:
|
||||
|
||||
await asyncio.gather(*[check_one(m) for m in models])
|
||||
return results
|
||||
|
||||
async def lookup_civitai_by_hash(self, file_hash: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Look up a model on CivitAI by its file hash.
|
||||
|
||||
Args:
|
||||
file_hash: SHA256 hash of the model file
|
||||
|
||||
Returns:
|
||||
Dict with model info including source_url, or None if not found
|
||||
"""
|
||||
if not file_hash:
|
||||
return None
|
||||
|
||||
try:
|
||||
session = await self._get_session()
|
||||
url = f"{self.CIVITAI_API_BASE}/model-versions/by-hash/{file_hash}"
|
||||
|
||||
async with session.get(url) as resp:
|
||||
if resp.status == 404:
|
||||
return None
|
||||
if resp.status != 200:
|
||||
print(f"CivitAI hash lookup failed: {resp.status}")
|
||||
return None
|
||||
|
||||
data = await resp.json()
|
||||
|
||||
# Build the source URL
|
||||
model_id = data.get("modelId")
|
||||
version_id = data.get("id")
|
||||
|
||||
if not model_id or not version_id:
|
||||
return None
|
||||
|
||||
source_url = f"https://civitai.com/models/{model_id}?modelVersionId={version_id}"
|
||||
|
||||
# Get file info
|
||||
files = data.get("files", [])
|
||||
primary_file = next(
|
||||
(f for f in files if f.get("primary", False)),
|
||||
files[0] if files else {}
|
||||
)
|
||||
|
||||
result = {
|
||||
"source_type": "civitai",
|
||||
"source_url": source_url,
|
||||
"model_id": model_id,
|
||||
"version_id": version_id,
|
||||
"model_name": data.get("model", {}).get("name", "Unknown"),
|
||||
"version_name": data.get("name"),
|
||||
"size": primary_file.get("sizeKB", 0) * 1024,
|
||||
"modified": data.get("publishedAt"),
|
||||
}
|
||||
|
||||
# Get hash info
|
||||
hashes = primary_file.get("hashes", {})
|
||||
if "SHA256" in hashes:
|
||||
result["remote_hash"] = hashes["SHA256"]
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error looking up hash on CivitAI: {e}")
|
||||
return None
|
||||
|
||||
async def batch_lookup_by_hash(
|
||||
self,
|
||||
models: List[Dict[str, Any]],
|
||||
max_concurrent: int = 3
|
||||
) -> Dict[int, Dict[str, Any]]:
|
||||
"""
|
||||
Look up multiple models by hash on CivitAI concurrently.
|
||||
|
||||
Args:
|
||||
models: List of model dicts with 'id' and 'file_hash' keys
|
||||
max_concurrent: Maximum concurrent requests (keep low to avoid rate limits)
|
||||
|
||||
Returns:
|
||||
Dict mapping model id to lookup result
|
||||
"""
|
||||
semaphore = asyncio.Semaphore(max_concurrent)
|
||||
results = {}
|
||||
|
||||
async def lookup_one(model: Dict[str, Any]):
|
||||
async with semaphore:
|
||||
file_hash = model.get("file_hash")
|
||||
if file_hash and not model.get("source_url"):
|
||||
# Small delay to avoid rate limiting
|
||||
await asyncio.sleep(0.5)
|
||||
info = await self.lookup_civitai_by_hash(file_hash)
|
||||
if info:
|
||||
results[model["id"]] = info
|
||||
|
||||
await asyncio.gather(*[lookup_one(m) for m in models])
|
||||
return results
|
||||
|
||||
@@ -71,6 +71,7 @@ class ModelSyncAPI:
|
||||
|
||||
app.router.add_post("/modelsync/check", self.check_updates)
|
||||
app.router.add_post("/modelsync/download/{id}", self.download_update)
|
||||
app.router.add_post("/modelsync/lookup-sources", self.lookup_sources)
|
||||
|
||||
app.router.add_get("/modelsync/config", self.get_config)
|
||||
app.router.add_post("/modelsync/config", self.set_config)
|
||||
@@ -270,6 +271,101 @@ class ModelSyncAPI:
|
||||
"model_id": model_id,
|
||||
})
|
||||
|
||||
async def lookup_sources(self, request: web.Request) -> web.Response:
|
||||
"""
|
||||
Look up source URL for ONE model at a time by querying CivitAI with file hash.
|
||||
Call repeatedly to process all models with progress feedback.
|
||||
"""
|
||||
data = await request.json() if request.body_exists else {}
|
||||
|
||||
# Get all models without source URLs
|
||||
all_models = self.store.get_all_models()
|
||||
models_to_lookup = [m for m in all_models if not m.get("source_url")]
|
||||
|
||||
if not models_to_lookup:
|
||||
return web.json_response({
|
||||
"status": "complete",
|
||||
"message": "All models already have sources",
|
||||
"remaining": 0,
|
||||
"total": len(all_models),
|
||||
})
|
||||
|
||||
# Find next model to process (prioritize ones with hashes)
|
||||
model = None
|
||||
for m in models_to_lookup:
|
||||
if m.get("file_hash"):
|
||||
model = m
|
||||
break
|
||||
|
||||
# If no model has hash, pick first one to hash
|
||||
if not model:
|
||||
model = models_to_lookup[0]
|
||||
|
||||
# Step 1: Compute hash if needed
|
||||
if not model.get("file_hash"):
|
||||
try:
|
||||
print(f"[ModelSync] Hashing {model['filename']}...")
|
||||
file_hash = compute_sha256_chunked(model["path"])
|
||||
self.store.update_model(model["id"], {"file_hash": file_hash})
|
||||
model["file_hash"] = file_hash
|
||||
return web.json_response({
|
||||
"status": "hashing",
|
||||
"message": f"Computed hash for {model['filename']}",
|
||||
"model": model["filename"],
|
||||
"remaining": len(models_to_lookup),
|
||||
"total": len(all_models),
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"[ModelSync] Error hashing {model['filename']}: {e}")
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": f"Failed to hash {model['filename']}: {str(e)}",
|
||||
"remaining": len(models_to_lookup),
|
||||
"total": len(all_models),
|
||||
})
|
||||
|
||||
# Step 2: Look up on CivitAI
|
||||
try:
|
||||
print(f"[ModelSync] Looking up {model['filename']} on CivitAI...")
|
||||
info = await self.tracker.lookup_civitai_by_hash(model["file_hash"])
|
||||
|
||||
if info:
|
||||
self.store.update_model(model["id"], {
|
||||
"source_url": info.get("source_url"),
|
||||
"source_type": info.get("source_type", "civitai"),
|
||||
"remote_hash": info.get("remote_hash"),
|
||||
})
|
||||
return web.json_response({
|
||||
"status": "found",
|
||||
"message": f"Found {model['filename']} on CivitAI: {info.get('model_name', 'Unknown')}",
|
||||
"model": model["filename"],
|
||||
"civitai_name": info.get("model_name"),
|
||||
"source_url": info.get("source_url"),
|
||||
"remaining": len(models_to_lookup) - 1,
|
||||
"total": len(all_models),
|
||||
})
|
||||
else:
|
||||
# Mark as checked but not found (use source_type to track)
|
||||
self.store.update_model(model["id"], {
|
||||
"source_type": "not_found",
|
||||
})
|
||||
return web.json_response({
|
||||
"status": "not_found",
|
||||
"message": f"{model['filename']} not found on CivitAI",
|
||||
"model": model["filename"],
|
||||
"remaining": len(models_to_lookup) - 1,
|
||||
"total": len(all_models),
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ModelSync] Error looking up {model['filename']}: {e}")
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": f"Lookup failed for {model['filename']}: {str(e)}",
|
||||
"remaining": len(models_to_lookup),
|
||||
"total": len(all_models),
|
||||
})
|
||||
|
||||
# Configuration endpoints
|
||||
|
||||
async def get_config(self, request: web.Request) -> web.Response:
|
||||
|
||||
Reference in New Issue
Block a user