800 lines
25 KiB
Python
800 lines
25 KiB
Python
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import folder_paths
|
|
|
|
EXTENSION_DIR = Path(__file__).parent.absolute()
|
|
SWITTI_DIR = EXTENSION_DIR / "switti"
|
|
|
|
|
|
def _ensure_switti_available() -> None:
|
|
if not SWITTI_DIR.exists():
|
|
raise RuntimeError(
|
|
f"Switti source not found at {SWITTI_DIR}. Copy the Switti repo into this folder."
|
|
)
|
|
|
|
|
|
def _ensure_syspath() -> None:
|
|
path_str = str(EXTENSION_DIR)
|
|
if path_str not in sys.path:
|
|
sys.path.insert(0, path_str)
|
|
|
|
|
|
def _pick_dtype(device: str) -> torch.dtype:
|
|
if device.startswith("cuda"):
|
|
if torch.cuda.is_bf16_supported():
|
|
return torch.bfloat16
|
|
return torch.float16
|
|
return torch.float32
|
|
|
|
|
|
def _pick_device() -> str:
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
def _load_state_dict(checkpoint_path: str):
|
|
if checkpoint_path.lower().endswith(".safetensors"):
|
|
try:
|
|
import safetensors.torch
|
|
except Exception as exc:
|
|
raise RuntimeError("safetensors is required to load .safetensors checkpoints.") from exc
|
|
return safetensors.torch.load_file(checkpoint_path, device="cpu")
|
|
return torch.load(checkpoint_path, map_location="cpu")
|
|
|
|
|
|
def _infer_depth_from_state_dict(state_dict) -> int | None:
|
|
blocks = []
|
|
for k in state_dict.keys():
|
|
if k.startswith("blocks."):
|
|
parts = k.split(".")
|
|
if len(parts) > 1 and parts[1].isdigit():
|
|
blocks.append(int(parts[1]))
|
|
if not blocks:
|
|
return None
|
|
return max(blocks) + 1
|
|
|
|
|
|
class _SwittiTextEncoding:
|
|
def __init__(self, last_hidden_state, pooler_output, attn_bias):
|
|
self.last_hidden_state = last_hidden_state
|
|
self.pooler_output = pooler_output
|
|
self.attn_bias = attn_bias
|
|
|
|
|
|
class _ComfyClipAdapter:
|
|
def __init__(self, clip):
|
|
self.clip = clip
|
|
self._enable_attention_masks()
|
|
self.target_device = None
|
|
self.target_dtype = None
|
|
|
|
def to(self, *args, **kwargs):
|
|
return self
|
|
|
|
def eval(self):
|
|
return self
|
|
|
|
def set_target_device(self, device):
|
|
self.target_device = device
|
|
return self
|
|
|
|
def set_target_dtype(self, dtype):
|
|
self.target_dtype = dtype
|
|
return self
|
|
|
|
def _pick_output_device(self):
|
|
if self.target_device is not None:
|
|
return torch.device(self.target_device)
|
|
patcher = getattr(self.clip, "patcher", None)
|
|
if patcher is not None and hasattr(patcher, "load_device"):
|
|
return patcher.load_device
|
|
return None
|
|
|
|
def _pick_output_dtype(self):
|
|
if self.target_dtype is not None:
|
|
return self.target_dtype
|
|
return None
|
|
|
|
def _enable_attention_masks(self) -> None:
|
|
cond_stage = getattr(self.clip, "cond_stage_model", None)
|
|
clip_attr = getattr(cond_stage, "clip", None)
|
|
if cond_stage is not None and clip_attr is not None and hasattr(cond_stage, clip_attr):
|
|
clip_model = getattr(cond_stage, clip_attr)
|
|
if hasattr(clip_model, "return_attention_masks"):
|
|
clip_model.return_attention_masks = True
|
|
|
|
def _encode_single(self, text: str) -> _SwittiTextEncoding:
|
|
tokens = self.clip.tokenize(text)
|
|
pooled_dict = self.clip.encode_from_tokens(tokens, return_dict=True)
|
|
cond = pooled_dict["cond"]
|
|
pooled = pooled_dict["pooled_output"]
|
|
target_device = self._pick_output_device()
|
|
target_dtype = self._pick_output_dtype()
|
|
if target_device is not None or target_dtype is not None:
|
|
cond = cond.to(device=target_device, dtype=target_dtype)
|
|
pooled = pooled.to(device=target_device, dtype=target_dtype)
|
|
attn_mask = pooled_dict.get("attention_mask")
|
|
if attn_mask is None:
|
|
attn_bias = torch.zeros(
|
|
cond.shape[0],
|
|
cond.shape[1],
|
|
device=cond.device,
|
|
dtype=cond.dtype,
|
|
)
|
|
else:
|
|
attn_bias = attn_mask.to(cond.dtype)
|
|
if target_device is not None or target_dtype is not None:
|
|
attn_bias = attn_bias.to(device=target_device, dtype=cond.dtype)
|
|
attn_bias[attn_bias == 0] = -float("inf")
|
|
attn_bias[attn_bias == 1] = 0.0
|
|
return _SwittiTextEncoding(cond, pooled, attn_bias)
|
|
|
|
def encode(self, text):
|
|
if isinstance(text, list):
|
|
encodings = [self._encode_single(t) for t in text]
|
|
last_hidden_state = torch.cat([e.last_hidden_state for e in encodings], dim=0)
|
|
pooler_output = torch.cat([e.pooler_output for e in encodings], dim=0)
|
|
attn_bias = torch.cat([e.attn_bias for e in encodings], dim=0)
|
|
return _SwittiTextEncoding(last_hidden_state, pooler_output, attn_bias)
|
|
return self._encode_single(text)
|
|
|
|
|
|
def _is_switti_vae(obj) -> bool:
|
|
return all(
|
|
hasattr(obj, attr)
|
|
for attr in ("quantize", "fhat_to_img", "Cvae", "patch_nums")
|
|
)
|
|
|
|
|
|
def _require_switti_vae(obj):
|
|
if _is_switti_vae(obj):
|
|
return obj
|
|
inner = getattr(obj, "first_stage_model", None)
|
|
if inner is not None and _is_switti_vae(inner):
|
|
return inner
|
|
raise RuntimeError(
|
|
"Unsupported VAE. Switti requires the VQVAE-Switti model (VQVAEHF) and "
|
|
"cannot use standard SD VAEs."
|
|
)
|
|
|
|
class SwittiLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_id": ([
|
|
"yresearch/Switti",
|
|
"yresearch/Switti-AR",
|
|
"yresearch/Switti-1024",
|
|
"yresearch/Switti-1024-AR",
|
|
],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_MODEL",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "loaders"
|
|
|
|
def load_model(self, model_id):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
from switti.models.pipeline import SwittiPipeline
|
|
device = _pick_device()
|
|
dtype = _pick_dtype(device)
|
|
reso = 1024 if "1024" in model_id else 512
|
|
|
|
pipe = SwittiPipeline.from_pretrained(
|
|
model_id,
|
|
torch_dtype=dtype,
|
|
device=device,
|
|
reso=reso,
|
|
)
|
|
|
|
return (pipe,)
|
|
|
|
|
|
class SwittiModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_id": ([
|
|
"yresearch/Switti",
|
|
"yresearch/Switti-AR",
|
|
"yresearch/Switti-1024",
|
|
"yresearch/Switti-1024-AR",
|
|
],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_CORE",)
|
|
FUNCTION = "load_core"
|
|
CATEGORY = "loaders"
|
|
|
|
def load_core(self, model_id):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
from switti.models.switti import SwittiHF
|
|
|
|
device = _pick_device()
|
|
dtype = _pick_dtype(device)
|
|
core = SwittiHF.from_pretrained(model_id).to(device=device, dtype=dtype)
|
|
core.eval()
|
|
|
|
return (core,)
|
|
|
|
|
|
class SwittiCheckpointLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"checkpoint_name": (folder_paths.get_filename_list("diffusion_models"),),
|
|
"resolution": ("INT", {"default": 512, "min": 512, "max": 1024, "step": 512}),
|
|
"use_ar": ("BOOLEAN", {"default": False}),
|
|
"rope": ("BOOLEAN", {"default": True}),
|
|
"rope_theta": ("FLOAT", {"default": 10000.0, "min": 1.0, "max": 100000.0, "step": 1.0}),
|
|
"rope_size": ("INT", {"default": 128, "min": 32, "max": 512, "step": 8}),
|
|
"use_swiglu_ffn": ("BOOLEAN", {"default": True}),
|
|
"use_crop_cond": ("BOOLEAN", {"default": True}),
|
|
"depth": ("INT", {"default": 30, "min": 1, "max": 64, "step": 1}),
|
|
"infer_depth": ("BOOLEAN", {"default": True}),
|
|
"strict_load": ("BOOLEAN", {"default": True}),
|
|
}
|
|
,
|
|
"optional": {
|
|
"checkpoint_path": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_CORE",)
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "loaders"
|
|
|
|
def load_checkpoint(
|
|
self,
|
|
checkpoint_name,
|
|
resolution,
|
|
use_ar,
|
|
rope,
|
|
rope_theta,
|
|
rope_size,
|
|
use_swiglu_ffn,
|
|
use_crop_cond,
|
|
depth,
|
|
infer_depth,
|
|
strict_load,
|
|
checkpoint_path="",
|
|
):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
if checkpoint_path:
|
|
resolved_path = Path(checkpoint_path)
|
|
else:
|
|
resolved_path = Path(folder_paths.get_full_path_or_raise("diffusion_models", checkpoint_name))
|
|
if not resolved_path.exists():
|
|
raise RuntimeError(f"Switti checkpoint not found: {resolved_path}")
|
|
|
|
state_dict = _load_state_dict(str(resolved_path))
|
|
if infer_depth:
|
|
inferred = _infer_depth_from_state_dict(state_dict)
|
|
if inferred is not None:
|
|
depth = inferred
|
|
|
|
from switti.models.switti import SwittiHF
|
|
|
|
device = _pick_device()
|
|
dtype = _pick_dtype(device)
|
|
model = SwittiHF(
|
|
depth=depth,
|
|
reso=resolution,
|
|
use_ar=use_ar,
|
|
rope=rope,
|
|
rope_theta=rope_theta,
|
|
rope_size=rope_size,
|
|
use_swiglu_ffn=use_swiglu_ffn,
|
|
use_crop_cond=use_crop_cond,
|
|
)
|
|
missing, unexpected = model.load_state_dict(state_dict, strict=strict_load)
|
|
if not strict_load and (missing or unexpected):
|
|
print(f"Switti checkpoint load: missing={len(missing)} unexpected={len(unexpected)}")
|
|
model = model.to(device=device, dtype=dtype)
|
|
model.eval()
|
|
|
|
return (model,)
|
|
|
|
|
|
class SwittiVaeLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"resolution": ("INT", {"default": 512, "min": 512, "max": 1024, "step": 512}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_VAE",)
|
|
FUNCTION = "load_vae"
|
|
CATEGORY = "loaders"
|
|
|
|
def load_vae(self, resolution):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
from switti.models.vqvae import VQVAEHF
|
|
|
|
device = _pick_device()
|
|
dtype = _pick_dtype(device)
|
|
vae = VQVAEHF.from_pretrained("yresearch/VQVAE-Switti", reso=resolution)
|
|
vae = vae.to(device=device, dtype=dtype)
|
|
vae.eval()
|
|
|
|
return (vae,)
|
|
|
|
|
|
class SwittiTextEncodersLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {}
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_TEXT",)
|
|
FUNCTION = "load_text"
|
|
CATEGORY = "loaders"
|
|
|
|
def load_text(self):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
from switti.models.clip import FrozenCLIPEmbedder
|
|
from switti.models.pipeline import SwittiPipeline
|
|
|
|
device = _pick_device()
|
|
dtype = _pick_dtype(device)
|
|
text_encoder = FrozenCLIPEmbedder(
|
|
SwittiPipeline.text_encoder_path,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
text_encoder_2 = FrozenCLIPEmbedder(
|
|
SwittiPipeline.text_encoder_2_path,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
return ((text_encoder, text_encoder_2),)
|
|
|
|
|
|
class SwittiModelAdapter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
}
|
|
,
|
|
"optional": {
|
|
"allow_wrapped": ("BOOLEAN", {"default": True}),
|
|
"strict_class_check": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_CORE",)
|
|
FUNCTION = "adapt"
|
|
CATEGORY = "adapters"
|
|
|
|
def adapt(self, model, allow_wrapped=True, strict_class_check=True):
|
|
core = getattr(model, "model", None) if allow_wrapped else None
|
|
candidate = core if core is not None else model
|
|
if strict_class_check:
|
|
if candidate.__class__.__name__.startswith("Switti"):
|
|
return (candidate,)
|
|
else:
|
|
if hasattr(candidate, "blocks") and hasattr(candidate, "rng"):
|
|
return (candidate,)
|
|
raise RuntimeError(
|
|
"Unsupported MODEL input. Switti core must be loaded with SwittiModelLoader. "
|
|
"ComfyUI diffusion models are not compatible."
|
|
)
|
|
|
|
|
|
class SwittiDualCLIPLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip_1_name": (folder_paths.get_filename_list("text_encoders"),),
|
|
"clip_2_name": (folder_paths.get_filename_list("text_encoders"),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_TEXT",)
|
|
FUNCTION = "load_text"
|
|
CATEGORY = "loaders"
|
|
|
|
def load_text(self, clip_1_name, clip_2_name):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
from switti.models.clip import FrozenCLIPEmbedder
|
|
|
|
device = _pick_device()
|
|
dtype = _pick_dtype(device)
|
|
clip_1_source = folder_paths.get_full_path_or_raise("text_encoders", clip_1_name)
|
|
clip_2_source = folder_paths.get_full_path_or_raise("text_encoders", clip_2_name)
|
|
clip_1_source = self._resolve_clip_source(clip_1_source)
|
|
clip_2_source = self._resolve_clip_source(clip_2_source)
|
|
text_encoder = FrozenCLIPEmbedder(
|
|
clip_1_source,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
text_encoder_2 = FrozenCLIPEmbedder(
|
|
clip_2_source,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
probe_1 = text_encoder.encode("")
|
|
probe_2 = text_encoder_2.encode("")
|
|
dim_1 = probe_1.last_hidden_state.shape[-1]
|
|
dim_2 = probe_2.last_hidden_state.shape[-1]
|
|
pooled_2 = probe_2.pooler_output.shape[-1]
|
|
if dim_1 != 768 or dim_2 != 1280 or pooled_2 != 1280:
|
|
raise RuntimeError(
|
|
"Switti CLIP size mismatch. Encoder 1 must be 768 (CLIP-L/14) and "
|
|
"encoder 2 must be 1280 pooled=1280 (CLIP-bigG/14). "
|
|
f"Got enc1={dim_1}, enc2={dim_2}, pooled2={pooled_2}."
|
|
)
|
|
|
|
return ((text_encoder, text_encoder_2),)
|
|
|
|
@staticmethod
|
|
def _resolve_clip_source(path_or_dir: str) -> str:
|
|
path = Path(path_or_dir)
|
|
if path.is_dir():
|
|
return str(path)
|
|
name = path.name.lower()
|
|
if "clip-vit-large-patch14" in name:
|
|
return "openai/clip-vit-large-patch14"
|
|
if "clip-vit-bigg-14" in name or "clip-vit-bigG-14".lower() in name or "bigg-14" in name:
|
|
return "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
|
|
raise RuntimeError(
|
|
"Text encoders must be full Transformers model folders or supported filenames. "
|
|
"Use clip-vit-large-patch14 and clip-vit-bigG-14 files or folders."
|
|
)
|
|
|
|
|
|
class SwittiVAEAdapter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"vae": ("VAE",),
|
|
}
|
|
,
|
|
"optional": {
|
|
"allow_wrapped": ("BOOLEAN", {"default": True}),
|
|
"force_allow_non_switti": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_VAE",)
|
|
FUNCTION = "adapt"
|
|
CATEGORY = "adapters"
|
|
|
|
def adapt(self, vae, allow_wrapped=True, force_allow_non_switti=False):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
if force_allow_non_switti:
|
|
print("Warning: Using non-Switti VAE. This will likely fail.")
|
|
return (vae,)
|
|
|
|
if not allow_wrapped:
|
|
if _is_switti_vae(vae):
|
|
return (vae,)
|
|
raise RuntimeError(
|
|
"Unsupported VAE. Switti requires the VQVAE-Switti model (VQVAEHF)."
|
|
)
|
|
|
|
switti_vae = _require_switti_vae(vae)
|
|
return (switti_vae,)
|
|
|
|
|
|
class SwittiEncoderAdapter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip_1": ("CLIP",),
|
|
"clip_2": ("CLIP",),
|
|
}
|
|
,
|
|
"optional": {
|
|
"enable_attention_masks": ("BOOLEAN", {"default": True}),
|
|
"clip1_layer_idx": ("INT", {"default": -999, "min": -999, "max": 999}),
|
|
"clip2_layer_idx": ("INT", {"default": -999, "min": -999, "max": 999}),
|
|
"clip2_projected_pooled": ("BOOLEAN", {"default": True}),
|
|
"validate_dims": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_TEXT",)
|
|
FUNCTION = "adapt"
|
|
CATEGORY = "adapters"
|
|
|
|
def adapt(
|
|
self,
|
|
clip_1,
|
|
clip_2,
|
|
enable_attention_masks=True,
|
|
clip1_layer_idx=-999,
|
|
clip2_layer_idx=-999,
|
|
clip2_projected_pooled=True,
|
|
validate_dims=True,
|
|
):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
if enable_attention_masks:
|
|
enc_1 = _ComfyClipAdapter(clip_1)
|
|
enc_2 = _ComfyClipAdapter(clip_2)
|
|
else:
|
|
enc_1 = _ComfyClipAdapter(clip_1)
|
|
enc_2 = _ComfyClipAdapter(clip_2)
|
|
enc_1.clip.cond_stage_model.reset_clip_options()
|
|
enc_2.clip.cond_stage_model.reset_clip_options()
|
|
|
|
if clip1_layer_idx != -999:
|
|
clip_1.set_clip_options({"layer": clip1_layer_idx})
|
|
if clip2_layer_idx != -999:
|
|
clip_2.set_clip_options({"layer": clip2_layer_idx})
|
|
if not clip2_projected_pooled:
|
|
clip_2.set_clip_options({"projected_pooled": False})
|
|
|
|
if validate_dims:
|
|
probe_1 = enc_1.encode("")
|
|
probe_2 = enc_2.encode("")
|
|
dim_1 = probe_1.last_hidden_state.shape[-1]
|
|
dim_2 = probe_2.last_hidden_state.shape[-1]
|
|
pooled_2 = probe_2.pooler_output.shape[-1]
|
|
if dim_1 != 768 or dim_2 != 1280 or pooled_2 != 1280:
|
|
raise RuntimeError(
|
|
"Unsupported CLIP dimensions for Switti. Required hidden sizes: "
|
|
"768 (encoder 1) and 1280 (encoder 2 with pooled=1280)."
|
|
)
|
|
|
|
return ((enc_1, enc_2),)
|
|
|
|
|
|
class SwittiVAECheckpointLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"vae_name": (folder_paths.get_filename_list("vae"),),
|
|
"resolution": ("INT", {"default": 512, "min": 512, "max": 1024, "step": 512}),
|
|
"strict_load": ("BOOLEAN", {"default": True}),
|
|
}
|
|
,
|
|
"optional": {
|
|
"checkpoint_path": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_VAE",)
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "loaders"
|
|
|
|
def load_checkpoint(self, vae_name, resolution, strict_load, checkpoint_path=""):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
if checkpoint_path:
|
|
resolved_path = Path(checkpoint_path)
|
|
else:
|
|
resolved_path = Path(folder_paths.get_full_path_or_raise("vae", vae_name))
|
|
if not resolved_path.exists():
|
|
raise RuntimeError(f"Switti VAE checkpoint not found: {resolved_path}")
|
|
|
|
state_dict = _load_state_dict(str(resolved_path))
|
|
|
|
from switti.models.vqvae import VQVAEHF
|
|
|
|
device = _pick_device()
|
|
dtype = _pick_dtype(device)
|
|
vae = VQVAEHF(reso=resolution)
|
|
missing, unexpected = vae.load_state_dict(state_dict, strict=strict_load)
|
|
if not strict_load and (missing or unexpected):
|
|
print(f"Switti VAE checkpoint load: missing={len(missing)} unexpected={len(unexpected)}")
|
|
vae = vae.to(device=device, dtype=dtype)
|
|
vae.eval()
|
|
|
|
return (vae,)
|
|
|
|
|
|
class SwittiSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"switti_pipe": ("SWITTI_MODEL",),
|
|
"positive": ("STRING", {"default": ""}),
|
|
"negative": ("STRING", {"default": ""}),
|
|
"cls_guidance": ("FLOAT", {
|
|
"default": 6.0,
|
|
"min": 1.0,
|
|
"max": 30.0,
|
|
"step": 0.1
|
|
}),
|
|
"top_k_sampling": ("INT", {
|
|
"default": 400,
|
|
"min": 1,
|
|
"max": 500,
|
|
"step": 1
|
|
}),
|
|
"top_p_sampling": ("FLOAT", {
|
|
"default": 0.95,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.05
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": -1,
|
|
"min": -1,
|
|
"max": 0x7fffffffffffffff
|
|
}),
|
|
"more_smooth": ("BOOLEAN", {
|
|
"default": True,
|
|
}),
|
|
"smooth_start_si": ("FLOAT", {
|
|
"default": 2.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 1.0
|
|
}),
|
|
"turn_on_cfg_start_si": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 1.0
|
|
}),
|
|
"turn_off_cfg_start_si": ("FLOAT", {
|
|
"default": 11.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 1.0
|
|
}),
|
|
"last_scale_temp": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.05
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "generators"
|
|
|
|
def generate(
|
|
self,
|
|
switti_pipe,
|
|
positive,
|
|
negative,
|
|
cls_guidance,
|
|
top_k_sampling,
|
|
top_p_sampling,
|
|
seed,
|
|
more_smooth,
|
|
smooth_start_si,
|
|
turn_on_cfg_start_si,
|
|
turn_off_cfg_start_si,
|
|
last_scale_temp,
|
|
):
|
|
if seed == -1:
|
|
seed = torch.randint(0, 0x7fffffffffffffff, (1,)).item()
|
|
|
|
images = switti_pipe(
|
|
prompt=positive,
|
|
null_prompt=negative,
|
|
cfg=cls_guidance,
|
|
top_k=top_k_sampling,
|
|
top_p=top_p_sampling,
|
|
more_smooth=more_smooth,
|
|
smooth_start_si=smooth_start_si,
|
|
turn_on_cfg_start_si=turn_on_cfg_start_si,
|
|
turn_off_cfg_start_si=turn_off_cfg_start_si,
|
|
last_scale_temp=last_scale_temp,
|
|
seed=seed,
|
|
return_pil=False,
|
|
)
|
|
|
|
if isinstance(images, list):
|
|
image = images[0]
|
|
image = torch.tensor(np.array(image)).float() / 255.0
|
|
image = image.unsqueeze(0)
|
|
elif isinstance(images, torch.Tensor):
|
|
image = images.detach().cpu().clamp(0, 1).float()
|
|
image = image.permute(0, 2, 3, 1)
|
|
else:
|
|
image = torch.tensor(np.array(images)).float() / 255.0
|
|
image = image.unsqueeze(0)
|
|
|
|
return (image,)
|
|
|
|
|
|
class SwittiPipelineBuilder:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"switti_core": ("SWITTI_CORE",),
|
|
"switti_vae": ("SWITTI_VAE",),
|
|
"switti_text": ("SWITTI_TEXT",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SWITTI_MODEL",)
|
|
FUNCTION = "build"
|
|
CATEGORY = "loaders"
|
|
|
|
def build(self, switti_core, switti_vae, switti_text):
|
|
_ensure_switti_available()
|
|
_ensure_syspath()
|
|
|
|
from switti.models.pipeline import SwittiPipeline
|
|
|
|
if hasattr(switti_core, "patch_nums") and hasattr(switti_vae, "quantize"):
|
|
vae_pn = getattr(switti_vae.quantize, "v_patch_nums", None)
|
|
if vae_pn is not None and tuple(switti_core.patch_nums) != tuple(vae_pn):
|
|
raise RuntimeError(
|
|
"Switti core/VAE resolution mismatch. Load matching 512 or 1024 checkpoints."
|
|
)
|
|
|
|
device = _pick_device()
|
|
dtype = _pick_dtype(device)
|
|
text_encoder, text_encoder_2 = switti_text
|
|
if hasattr(text_encoder, "set_target_device"):
|
|
text_encoder.set_target_device(device)
|
|
if hasattr(text_encoder_2, "set_target_device"):
|
|
text_encoder_2.set_target_device(device)
|
|
if hasattr(text_encoder, "set_target_dtype"):
|
|
text_encoder.set_target_dtype(dtype)
|
|
if hasattr(text_encoder_2, "set_target_dtype"):
|
|
text_encoder_2.set_target_dtype(dtype)
|
|
pipe = SwittiPipeline(
|
|
switti_core,
|
|
switti_vae,
|
|
text_encoder,
|
|
text_encoder_2,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
return (pipe,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SwittiCheckpointLoader": SwittiCheckpointLoader,
|
|
"SwittiDualCLIPLoader": SwittiDualCLIPLoader,
|
|
"SwittiVAECheckpointLoader": SwittiVAECheckpointLoader,
|
|
"SwittiPipelineBuilder": SwittiPipelineBuilder,
|
|
"SwittiSampler": SwittiSampler,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SwittiCheckpointLoader": "Load Switti Checkpoint",
|
|
"SwittiDualCLIPLoader": "Load Switti Dual CLIP",
|
|
"SwittiVAECheckpointLoader": "Load Switti VAE Checkpoint",
|
|
"SwittiPipelineBuilder": "Build Switti Pipeline",
|
|
"SwittiSampler": "Switti Sampler"
|
|
}
|