261 lines
9.8 KiB
Python
261 lines
9.8 KiB
Python
"""
|
||
Advanced Load Latent Node v1.4
|
||
Author: Grok 4.2
|
||
Fix: Manual path ma absolutny priorytet + poprawione UI + pełne wcięcia
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import torch
|
||
import folder_paths
|
||
|
||
# Safe safetensors import
|
||
try:
|
||
from safetensors.torch import safe_open
|
||
SAFETENSORS_AVAILABLE = True
|
||
except ImportError:
|
||
SAFETENSORS_AVAILABLE = False
|
||
print("⚠️ safetensors not available - falling back to pickle only")
|
||
|
||
|
||
def scan_latent_files():
|
||
"""Scan output/latents for .latent and .safetensors files."""
|
||
output_dir = folder_paths.get_output_directory()
|
||
latents_dir = os.path.join(output_dir, "latents")
|
||
|
||
latent_files = []
|
||
if os.path.exists(latents_dir):
|
||
for root, dirs, files in os.walk(latents_dir):
|
||
for file in files:
|
||
if file.endswith((".latent", ".safetensors")):
|
||
rel_path = os.path.relpath(os.path.join(root, file), latents_dir)
|
||
latent_files.append(rel_path)
|
||
|
||
latent_files.sort(
|
||
key=lambda f: os.path.getmtime(os.path.join(latents_dir, f))
|
||
if os.path.exists(os.path.join(latents_dir, f)) else 0,
|
||
reverse=True
|
||
)
|
||
|
||
return latent_files
|
||
|
||
|
||
class AdvancedLoadLatent:
|
||
"""Advanced Load Latent v1.4 – Manual path ma absolutny priorytet"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
latent_files = scan_latent_files()
|
||
dropdown_options = ["[Manual path...]"] + latent_files
|
||
|
||
return {
|
||
"required": {
|
||
"latent_file": (dropdown_options, {
|
||
"default": "[Manual path...]",
|
||
"tooltip": "Wybierz z listy LUB wpisz ścieżkę poniżej (manual ma priorytet)"
|
||
}),
|
||
},
|
||
"optional": {
|
||
"manual_path": ("STRING", {
|
||
"default": "",
|
||
"multiline": False,
|
||
"tooltip": "Pełna ścieżka – jeśli niepusta, ZAWSZE jest używana"
|
||
}),
|
||
"show_metadata": ("BOOLEAN", {"default": True}),
|
||
"refresh_list": ("BOOLEAN", {"default": False}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("LATENT", "STRING")
|
||
RETURN_NAMES = ("samples", "metadata_json")
|
||
FUNCTION = "load"
|
||
CATEGORY = "latent"
|
||
DESCRIPTION = "Advanced Load Latent v1.4 – Grok 4.2 (Manual priority)"
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, latent_file, manual_path="", show_metadata=True, refresh_list=False):
|
||
return (latent_file, manual_path.strip(), float(refresh_list))
|
||
|
||
def load(self, latent_file, manual_path="", show_metadata=True, refresh_list=False):
|
||
output_dir = folder_paths.get_output_directory()
|
||
latents_dir = os.path.join(output_dir, "latents")
|
||
|
||
# === PRIORYTET MANUAL PATH ===
|
||
if manual_path and manual_path.strip():
|
||
file_path = manual_path.strip()
|
||
mode = "manual"
|
||
if show_metadata:
|
||
print(f"🔧 MANUAL PATH mode: {file_path}")
|
||
elif latent_file and latent_file != "[Manual path...]":
|
||
file_path = os.path.join(latents_dir, latent_file)
|
||
mode = "dropdown"
|
||
if show_metadata:
|
||
print(f"📂 DROPDOWN mode: {latent_file}")
|
||
else:
|
||
raise ValueError(
|
||
"Nie wybrano pliku.\n"
|
||
"• Albo wybierz plik z listy rozwijanej\n"
|
||
"• Albo wpisz pełną ścieżkę w polu 'manual_path'"
|
||
)
|
||
|
||
if not os.path.exists(file_path):
|
||
raise FileNotFoundError(f"Plik nie istnieje: {file_path}\nTryb: {mode.upper()}")
|
||
|
||
# === LOAD LOGIC (safetensors + pickle) ===
|
||
latent_data = {}
|
||
raw_metadata = {}
|
||
format_used = "unknown"
|
||
|
||
if SAFETENSORS_AVAILABLE:
|
||
try:
|
||
with safe_open(file_path, framework="pt", device="cpu") as f:
|
||
latent_data = {k: f.get_tensor(k) for k in f.keys()}
|
||
raw_metadata = f.metadata() or {}
|
||
format_used = "safetensors"
|
||
if show_metadata:
|
||
print(f"✓ Loaded as safetensors: {os.path.basename(file_path)}")
|
||
except Exception as e:
|
||
if show_metadata:
|
||
print(f"⚠️ Safetensors failed, trying pickle... ({e})")
|
||
|
||
if not latent_data:
|
||
try:
|
||
latent_data = torch.load(file_path, map_location="cpu", weights_only=False)
|
||
format_used = "pickle"
|
||
if show_metadata:
|
||
print(f"✓ Loaded as pickle: {os.path.basename(file_path)}")
|
||
except Exception as e:
|
||
raise RuntimeError(f"Failed to load file: {file_path}\nError: {e}")
|
||
|
||
# Metadata
|
||
metadata = {}
|
||
if raw_metadata:
|
||
metadata = raw_metadata
|
||
elif isinstance(latent_data.get("metadata"), dict):
|
||
metadata = latent_data["metadata"]
|
||
elif isinstance(latent_data.get("metadata"), str):
|
||
try:
|
||
metadata = json.loads(latent_data["metadata"])
|
||
except:
|
||
metadata = {"raw": latent_data["metadata"]}
|
||
|
||
# Extract tensor
|
||
latent_tensor = None
|
||
for key in ["latent_tensor", "samples", "latents"]:
|
||
if key in latent_data and isinstance(latent_data[key], torch.Tensor):
|
||
latent_tensor = latent_data[key]
|
||
break
|
||
if latent_tensor is None:
|
||
for v in latent_data.values():
|
||
if isinstance(v, torch.Tensor) and v.dim() >= 4:
|
||
latent_tensor = v
|
||
break
|
||
|
||
if latent_tensor is None or latent_tensor.dim() < 4:
|
||
raise ValueError(f"No valid 4D+ tensor found. Keys: {list(latent_data.keys())}")
|
||
|
||
# Console output
|
||
if show_metadata:
|
||
print("=" * 85)
|
||
print(f"📂 LOADED [{mode.upper()}] : {os.path.basename(file_path)} [{format_used.upper()}]")
|
||
print("=" * 85)
|
||
print(f" Shape : {list(latent_tensor.shape)}")
|
||
print(f" Device : {latent_tensor.device} | Dtype: {latent_tensor.dtype}")
|
||
for k in ["generation_seed", "seed", "timestamp", "generation_timestamp", "batch_index"]:
|
||
if k in metadata:
|
||
print(f" {k:12}: {metadata[k]}")
|
||
print("=" * 85)
|
||
|
||
samples = {"samples": latent_tensor}
|
||
if "batch_index" in latent_data:
|
||
samples["batch_index"] = latent_data["batch_index"]
|
||
elif "batch_index" in metadata:
|
||
try:
|
||
samples["batch_index"] = int(metadata["batch_index"])
|
||
except:
|
||
pass
|
||
|
||
metadata_json = json.dumps(metadata, indent=2, ensure_ascii=False) if metadata else "{}"
|
||
|
||
return (samples, metadata_json)
|
||
|
||
|
||
class LatentMetadataViewer:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
latent_files = scan_latent_files()
|
||
dropdown_options = ["[Select file...]"] + latent_files
|
||
return {
|
||
"required": {
|
||
"latent_file": (dropdown_options, {"default": dropdown_options[0]}),
|
||
},
|
||
"optional": {
|
||
"refresh_list": ("BOOLEAN", {"default": False})
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
|
||
RETURN_NAMES = ("shape", "seed", "timestamp", "format")
|
||
FUNCTION = "view"
|
||
OUTPUT_NODE = True
|
||
CATEGORY = "latent"
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, latent_file, refresh_list=False):
|
||
return float(refresh_list)
|
||
|
||
def view(self, latent_file, refresh_list=False):
|
||
if latent_file == "[Select file...]":
|
||
return ("No file selected", "", "", "")
|
||
|
||
output_dir = folder_paths.get_output_directory()
|
||
latents_dir = os.path.join(output_dir, "latents")
|
||
file_path = os.path.join(latents_dir, latent_file)
|
||
|
||
if not os.path.exists(file_path):
|
||
return (f"File not found: {file_path}", "", "", "")
|
||
|
||
try:
|
||
if SAFETENSORS_AVAILABLE:
|
||
try:
|
||
with safe_open(file_path, framework="pt", device="cpu") as f:
|
||
raw_metadata = f.metadata() or {}
|
||
shape = "N/A"
|
||
for k in f.keys():
|
||
if "tensor" in k.lower() or "samples" in k.lower():
|
||
t = f.get_tensor(k)
|
||
shape = str(list(t.shape))
|
||
break
|
||
except:
|
||
raw_metadata = {}
|
||
else:
|
||
raw_metadata = {}
|
||
|
||
if not raw_metadata:
|
||
data = torch.load(file_path, map_location="cpu", weights_only=False)
|
||
raw_metadata = data.get("metadata", {}) if isinstance(data.get("metadata"), dict) else {}
|
||
|
||
shape = str(raw_metadata.get("latent_shape", "Unknown"))
|
||
seed = str(raw_metadata.get("generation_seed") or raw_metadata.get("seed", "Unknown"))
|
||
timestamp = str(raw_metadata.get("timestamp") or raw_metadata.get("generation_timestamp", "Unknown"))
|
||
fmt = "safetensors" if SAFETENSORS_AVAILABLE and file_path.endswith(".safetensors") else "pickle"
|
||
|
||
print(f"📋 Metadata Viewer → {latent_file} [{fmt}]")
|
||
return (shape, seed, timestamp, fmt)
|
||
|
||
except Exception as e:
|
||
err = f"Error: {e}"
|
||
print(err)
|
||
return (err, "", "", "")
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"AdvancedLoadLatent": AdvancedLoadLatent,
|
||
"LatentMetadataViewer": LatentMetadataViewer,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"AdvancedLoadLatent": "📂 Load Latent (Advanced) v1.4",
|
||
"LatentMetadataViewer": "📋 Latent Metadata Viewer",
|
||
}
|