292 lines
10 KiB
Python
292 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
import inspect
|
|
|
|
try:
|
|
import torch
|
|
except Exception as exc:
|
|
torch = None
|
|
_TORCH_IMPORT_ERROR = exc
|
|
else:
|
|
_TORCH_IMPORT_ERROR = None
|
|
|
|
try:
|
|
import folder_paths
|
|
except Exception as exc:
|
|
folder_paths = None
|
|
_FOLDER_PATHS_IMPORT_ERROR = exc
|
|
else:
|
|
_FOLDER_PATHS_IMPORT_ERROR = None
|
|
|
|
try:
|
|
import comfy.model_management as model_management
|
|
except Exception as exc:
|
|
model_management = None
|
|
_MODEL_MANAGEMENT_IMPORT_ERROR = exc
|
|
else:
|
|
_MODEL_MANAGEMENT_IMPORT_ERROR = None
|
|
|
|
try:
|
|
import comfy.sd as comfy_sd
|
|
except Exception as exc:
|
|
comfy_sd = None
|
|
_COMFY_SD_IMPORT_ERROR = exc
|
|
else:
|
|
_COMFY_SD_IMPORT_ERROR = None
|
|
|
|
from ..sampling.nodes import _common_ksampler
|
|
from .request import LLSUniversalGenerationRequest
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class LoadedModels:
|
|
model: Any
|
|
clip: Any
|
|
vae: Any
|
|
|
|
|
|
class UniversalBackendBase(ABC):
|
|
"""统一 backend 接口。"""
|
|
|
|
family_name = "BASE"
|
|
|
|
@abstractmethod
|
|
def load_models(self, request: LLSUniversalGenerationRequest) -> LoadedModels:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def encode_prompt(self, request: LLSUniversalGenerationRequest, clip: Any) -> tuple[Any, Any]:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def prepare_latent(self, request: LLSUniversalGenerationRequest, model: Any) -> dict[str, Any]:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def sample(
|
|
self,
|
|
request: LLSUniversalGenerationRequest,
|
|
model: Any,
|
|
positive: Any,
|
|
negative: Any,
|
|
latent: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def decode(self, request: LLSUniversalGenerationRequest, vae: Any, samples: dict[str, Any]) -> Any:
|
|
raise NotImplementedError
|
|
|
|
def generate(self, request: LLSUniversalGenerationRequest) -> Any:
|
|
request = request.validate()
|
|
loaded = self.load_models(request)
|
|
positive, negative = self.encode_prompt(request, loaded.clip)
|
|
latent = self.prepare_latent(request, loaded.model)
|
|
sampled = self.sample(request, loaded.model, positive, negative, latent)
|
|
return self.decode(request, loaded.vae, sampled)
|
|
|
|
# ---------- 共用辅助 ----------
|
|
|
|
def _require_torch(self) -> None:
|
|
if torch is None:
|
|
raise RuntimeError(
|
|
"[LLS] PyTorch is not available. "
|
|
"Run this node inside ComfyUI's Python environment."
|
|
) from _TORCH_IMPORT_ERROR
|
|
|
|
def _require_comfy_runtime(self) -> None:
|
|
if folder_paths is None:
|
|
raise RuntimeError(
|
|
"[LLS] folder_paths is not available. "
|
|
"Place this plugin under ComfyUI/custom_nodes/ and restart ComfyUI."
|
|
) from _FOLDER_PATHS_IMPORT_ERROR
|
|
if comfy_sd is None:
|
|
raise RuntimeError(
|
|
"[LLS] comfy.sd is not available. "
|
|
"Run this node inside a ComfyUI environment."
|
|
) from _COMFY_SD_IMPORT_ERROR
|
|
if model_management is None:
|
|
raise RuntimeError(
|
|
"[LLS] comfy.model_management is not available. "
|
|
"Run this node inside a ComfyUI environment."
|
|
) from _MODEL_MANAGEMENT_IMPORT_ERROR
|
|
|
|
def _resolve_checkpoint_path(self, model_name: str) -> str:
|
|
self._require_comfy_runtime()
|
|
try:
|
|
return folder_paths.get_full_path_or_raise("checkpoints", model_name)
|
|
except (AttributeError, TypeError):
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", model_name)
|
|
if ckpt_path:
|
|
return ckpt_path
|
|
except FileNotFoundError as exc:
|
|
raise RuntimeError(
|
|
f"[LLS] Checkpoint '{model_name}' was not found in ComfyUI checkpoints."
|
|
) from exc
|
|
raise RuntimeError(
|
|
f"[LLS] Checkpoint '{model_name}' was not found in ComfyUI checkpoints."
|
|
)
|
|
|
|
def _load_checkpoint_bundle(self, request: LLSUniversalGenerationRequest) -> LoadedModels:
|
|
ckpt_path = self._resolve_checkpoint_path(request.model_name)
|
|
try:
|
|
embedding_directory = folder_paths.get_folder_paths("embeddings")
|
|
except Exception:
|
|
embedding_directory = []
|
|
|
|
try:
|
|
model, clip, vae = comfy_sd.load_checkpoint_guess_config(
|
|
ckpt_path,
|
|
output_vae=True,
|
|
output_clip=True,
|
|
embedding_directory=embedding_directory,
|
|
)[:3]
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
f"[LLS] Failed to load checkpoint '{request.model_name}' for {self.family_name}: {exc}"
|
|
) from exc
|
|
|
|
if model is None or clip is None or vae is None:
|
|
raise RuntimeError(
|
|
f"[LLS] Checkpoint '{request.model_name}' did not produce MODEL/CLIP/VAE required by {self.family_name}."
|
|
)
|
|
|
|
return LoadedModels(model=model, clip=clip, vae=vae)
|
|
|
|
def _infer_loaded_family(self, model: Any) -> str:
|
|
model_impl = getattr(model, "model", model)
|
|
latent_format = getattr(model_impl, "latent_format", None)
|
|
latent_name = type(latent_format).__name__
|
|
config_name = type(getattr(model_impl, "model_config", None)).__name__
|
|
combined = f"{latent_name} {config_name}".lower()
|
|
|
|
if "flux" in combined:
|
|
return "FLUX"
|
|
if "sdxl" in combined:
|
|
return "SDXL"
|
|
if "sd15" in combined:
|
|
return "SD1.5"
|
|
return "UNKNOWN"
|
|
|
|
def _validate_loaded_family(self, model: Any, expected: str) -> None:
|
|
loaded_family = self._infer_loaded_family(model)
|
|
if loaded_family == "UNKNOWN":
|
|
return
|
|
if loaded_family != expected:
|
|
raise RuntimeError(
|
|
f"[LLS] Selected backend '{expected}' does not match loaded checkpoint family '{loaded_family}'."
|
|
)
|
|
|
|
def _get_latent_spec(self, model: Any) -> tuple[int, int]:
|
|
model_impl = getattr(model, "model", model)
|
|
latent_format = getattr(model_impl, "latent_format", None)
|
|
channels = int(getattr(latent_format, "latent_channels", 4))
|
|
downscale = int(getattr(latent_format, "spacial_downscale_ratio", 8))
|
|
return channels, downscale
|
|
|
|
def _prepare_empty_latent(
|
|
self,
|
|
request: LLSUniversalGenerationRequest,
|
|
model: Any,
|
|
) -> dict[str, Any]:
|
|
self._require_torch()
|
|
self._require_comfy_runtime()
|
|
|
|
latent_channels, downscale = self._get_latent_spec(model)
|
|
aligned_request, changed = request.align_to(downscale)
|
|
if changed:
|
|
print(
|
|
f"[LLS] {self.family_name}: adjusted size "
|
|
f"{request.width}x{request.height} -> {aligned_request.width}x{aligned_request.height} "
|
|
f"to match latent downscale ratio {downscale}."
|
|
)
|
|
|
|
latent = torch.zeros(
|
|
[1, latent_channels, aligned_request.height // downscale, aligned_request.width // downscale],
|
|
device=model_management.intermediate_device(),
|
|
dtype=model_management.intermediate_dtype(),
|
|
)
|
|
return {"samples": latent, "downscale_ratio_spacial": downscale}
|
|
|
|
def _sample_latent(
|
|
self,
|
|
request: LLSUniversalGenerationRequest,
|
|
model: Any,
|
|
positive: Any,
|
|
negative: Any,
|
|
latent: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
actual_seed = request.resolve_seed()
|
|
try:
|
|
return _common_ksampler(
|
|
model=model,
|
|
seed=actual_seed,
|
|
steps=request.steps,
|
|
cfg=request.cfg,
|
|
sampler_name=request.sampler_name,
|
|
scheduler=request.scheduler,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=latent,
|
|
denoise=request.denoise,
|
|
)
|
|
except Exception as exc:
|
|
raise RuntimeError(
|
|
f"[LLS] {self.family_name} sampling failed for checkpoint '{request.model_name}': {exc}"
|
|
) from exc
|
|
|
|
def _decode_latent(self, vae: Any, samples: dict[str, Any]) -> Any:
|
|
latent = samples["samples"]
|
|
if getattr(latent, "is_nested", False):
|
|
latent = latent.unbind()[0]
|
|
images = vae.decode(latent)
|
|
if len(getattr(images, "shape", ())) == 5:
|
|
images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
|
|
return images
|
|
|
|
def _encode_tokens(self, clip: Any, tokens: Any, add_dict: dict[str, Any] | None = None) -> Any:
|
|
add_dict = dict(add_dict or {})
|
|
|
|
scheduled = getattr(clip, "encode_from_tokens_scheduled", None)
|
|
if callable(scheduled):
|
|
return scheduled(tokens, add_dict=add_dict)
|
|
|
|
plain_encode = getattr(clip, "encode_from_tokens", None)
|
|
if not callable(plain_encode):
|
|
raise RuntimeError(
|
|
f"[LLS] {self.family_name} backend requires CLIP.encode_from_tokens_scheduled() "
|
|
"or CLIP.encode_from_tokens()."
|
|
)
|
|
|
|
try:
|
|
signature = inspect.signature(plain_encode)
|
|
except (TypeError, ValueError):
|
|
signature = None
|
|
|
|
supports_return_dict = signature is not None and "return_dict" in signature.parameters
|
|
if supports_return_dict:
|
|
pooled_dict = plain_encode(tokens, return_pooled=True, return_dict=True)
|
|
else:
|
|
encoded = plain_encode(tokens, return_pooled=True)
|
|
if isinstance(encoded, tuple):
|
|
cond, pooled = encoded[:2]
|
|
pooled_dict = {"cond": cond, "pooled_output": pooled}
|
|
else:
|
|
pooled_dict = {"cond": encoded}
|
|
|
|
if not isinstance(pooled_dict, dict) or "cond" not in pooled_dict:
|
|
raise RuntimeError(
|
|
f"[LLS] {self.family_name} backend received unsupported CLIP encode output."
|
|
)
|
|
|
|
cond = pooled_dict.pop("cond")
|
|
pooled_dict.update(add_dict)
|
|
return [[cond, pooled_dict]]
|
|
|
|
def _encode_standard_prompt(self, clip: Any, prompt: str, add_dict: dict[str, Any] | None = None) -> Any:
|
|
tokens = clip.tokenize(prompt)
|
|
return self._encode_tokens(clip, tokens, add_dict=add_dict)
|