145 lines
4.9 KiB
Python
145 lines
4.9 KiB
Python
import os
|
|
import re
|
|
from dataclasses import dataclass
|
|
|
|
import safetensors.torch
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
|
|
@dataclass
|
|
class EmbeddingBundle:
|
|
clip_l: torch.Tensor | None
|
|
clip_g: torch.Tensor | None
|
|
metadata: dict[str, str]
|
|
source_path: str
|
|
|
|
@property
|
|
def token_count(self) -> int:
|
|
if self.clip_l is not None:
|
|
return int(self.clip_l.shape[0])
|
|
if self.clip_g is not None:
|
|
return int(self.clip_g.shape[0])
|
|
return 0
|
|
|
|
@property
|
|
def is_sdxl(self) -> bool:
|
|
return self.clip_l is not None and self.clip_g is not None
|
|
|
|
|
|
def load_embedding_bundle(embedding_path: str) -> EmbeddingBundle:
|
|
if embedding_path.lower().endswith((".safetensors", ".sft")):
|
|
return _load_safetensors_embedding(embedding_path)
|
|
return _load_pt_embedding(embedding_path)
|
|
|
|
|
|
def normalize_output_name(output_name: str, fallback_stem: str) -> str:
|
|
name = (output_name or "").strip()
|
|
if not name:
|
|
name = fallback_stem
|
|
stem = os.path.splitext(os.path.basename(name))[0]
|
|
return normalize_output_component(stem, fallback_stem)
|
|
|
|
|
|
def normalize_output_component(name: str, fallback_component: str) -> str:
|
|
component = re.sub(r"[^A-Za-z0-9._-]+", "_", (name or "").strip()).strip("._")
|
|
return component or fallback_component
|
|
|
|
|
|
def build_output_path(output_dir: str, output_name: str, fallback_stem: str, overwrite: bool) -> tuple[str, str]:
|
|
embedding_name = normalize_output_name(output_name, fallback_stem)
|
|
output_path = os.path.join(output_dir, f"{embedding_name}.safetensors")
|
|
if os.path.exists(output_path) and not overwrite:
|
|
raise FileExistsError(f"Output embedding already exists: {output_path}")
|
|
return output_path, embedding_name
|
|
|
|
|
|
def save_sdxl_embedding(
|
|
output_path: str,
|
|
clip_l: torch.Tensor,
|
|
clip_g: torch.Tensor,
|
|
metadata: dict[str, str],
|
|
) -> None:
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
safetensors.torch.save_file(
|
|
{"clip_l": clip_l.contiguous().cpu(), "clip_g": clip_g.contiguous().cpu()},
|
|
output_path,
|
|
metadata=metadata,
|
|
)
|
|
|
|
|
|
def _load_safetensors_embedding(embedding_path: str) -> EmbeddingBundle:
|
|
tensors = {}
|
|
with safe_open(embedding_path, framework="pt", device="cpu") as handle:
|
|
metadata = handle.metadata() or {}
|
|
for key in handle.keys():
|
|
tensors[key] = handle.get_tensor(key).clone()
|
|
|
|
clip_l = _coerce_tensor(tensors.get("clip_l"))
|
|
clip_g = _coerce_tensor(tensors.get("clip_g"))
|
|
if clip_l is None and clip_g is None:
|
|
inferred_key, inferred_tensor = _infer_single_tensor(tensors)
|
|
if inferred_key == "clip_l":
|
|
clip_l = inferred_tensor
|
|
elif inferred_key == "clip_g":
|
|
clip_g = inferred_tensor
|
|
|
|
return EmbeddingBundle(clip_l=clip_l, clip_g=clip_g, metadata=dict(metadata), source_path=embedding_path)
|
|
|
|
|
|
def _load_pt_embedding(embedding_path: str) -> EmbeddingBundle:
|
|
import comfy.utils
|
|
|
|
payload = comfy.utils.load_torch_file(embedding_path, safe_load=True, device=torch.device("cpu"))
|
|
metadata = {}
|
|
clip_l = None
|
|
clip_g = None
|
|
|
|
if isinstance(payload, dict):
|
|
for key in ("name", "step", "sd_checkpoint", "sd_checkpoint_name"):
|
|
value = payload.get(key)
|
|
if value is not None:
|
|
metadata[f"legacy.{key}"] = str(value)
|
|
|
|
if "clip_l" in payload:
|
|
clip_l = _coerce_tensor(payload["clip_l"])
|
|
if "clip_g" in payload:
|
|
clip_g = _coerce_tensor(payload["clip_g"])
|
|
|
|
string_to_param = payload.get("string_to_param")
|
|
if isinstance(string_to_param, dict) and (clip_l is None and clip_g is None):
|
|
first_tensor = next(iter(string_to_param.values()), None)
|
|
inferred = _coerce_tensor(first_tensor)
|
|
if inferred is not None:
|
|
if inferred.shape[-1] == 768:
|
|
clip_l = inferred
|
|
elif inferred.shape[-1] == 1280:
|
|
clip_g = inferred
|
|
|
|
if clip_l is None and clip_g is None:
|
|
raise RuntimeError(f"Unsupported embedding format: {embedding_path}")
|
|
|
|
return EmbeddingBundle(clip_l=clip_l, clip_g=clip_g, metadata=metadata, source_path=embedding_path)
|
|
|
|
|
|
def _infer_single_tensor(tensors: dict[str, torch.Tensor]) -> tuple[str | None, torch.Tensor | None]:
|
|
for tensor in tensors.values():
|
|
inferred = _coerce_tensor(tensor)
|
|
if inferred is None:
|
|
continue
|
|
if inferred.shape[-1] == 768:
|
|
return "clip_l", inferred
|
|
if inferred.shape[-1] == 1280:
|
|
return "clip_g", inferred
|
|
return None, None
|
|
|
|
|
|
def _coerce_tensor(tensor: torch.Tensor | None) -> torch.Tensor | None:
|
|
if tensor is None or not torch.is_tensor(tensor):
|
|
return None
|
|
if tensor.ndim == 1:
|
|
return tensor.view(1, -1)
|
|
if tensor.ndim >= 2:
|
|
return tensor.reshape(-1, tensor.shape[-1])
|
|
return None
|