951 lines
33 KiB
Python
951 lines
33 KiB
Python
"""
|
|
STARFlow ComfyUI Nodes
|
|
|
|
Implements custom nodes for STARFlow text-to-image generation.
|
|
"""
|
|
|
|
import os
|
|
import math
|
|
from types import SimpleNamespace
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import folder_paths
|
|
import comfy.model_management as mm
|
|
from transformers import AutoConfig, T5Tokenizer, T5EncoderModel
|
|
from safetensors.torch import load_file
|
|
|
|
def _nearest_multiple(x: float, base: int = 8) -> int:
|
|
return int(round(x / base)) * base
|
|
|
|
|
|
def aspect_ratio_to_image_size(target_size, ratio, multiple=8):
|
|
if ratio is None:
|
|
return target_size, target_size
|
|
if isinstance(ratio, str):
|
|
rw, rh = map(int, ratio.split(":"))
|
|
ratio = rw / rh
|
|
area = target_size ** 2
|
|
out_h = _nearest_multiple(math.sqrt(area / ratio), multiple)
|
|
out_w = _nearest_multiple(math.sqrt(area * ratio), multiple)
|
|
return out_h, out_w
|
|
|
|
|
|
def preprocess_text(text, use_template=False, aspect_ratio=None, fps=None, noise_std=None):
|
|
modes = ["an image"] * len(text)
|
|
if fps is not None:
|
|
if isinstance(fps, torch.Tensor):
|
|
fps = [int(f) for f in fps.tolist()]
|
|
elif isinstance(fps, int):
|
|
fps = [fps] * len(text)
|
|
modes = ["a video" if f > 0 else "an image" for f in fps]
|
|
text = [f"A video with {f} fps:\n{txt}\n" if f > 0 else f"An image:\n{txt}\n"
|
|
for txt, f in zip(text, fps)]
|
|
if noise_std is not None:
|
|
if isinstance(noise_std, torch.Tensor):
|
|
noise_std = [int(n * 1000) for n in noise_std.view(-1).tolist()]
|
|
elif isinstance(noise_std, float):
|
|
noise_std = [int(noise_std * 1000)] * len(text)
|
|
text = [f"Noise Level {n}:\n{txt}" for n, txt in zip(noise_std, text)]
|
|
if aspect_ratio is not None:
|
|
text = [f"{txt}\n in a {aspect_ratio} aspect ratio.\n" for txt in text]
|
|
if use_template:
|
|
template = "<start_of_turn>user\nPlease generate {mode} about: {prompt}<end_of_turn>\n<start_of_turn>model\n"
|
|
text = [template.format(prompt=txt, mode=mode) for txt, mode in zip(text, modes)]
|
|
return text
|
|
|
|
|
|
@torch.no_grad()
|
|
def encode_text(text_encoder, tokenizer, text, max_length, device, return_tokens=False, **kwargs):
|
|
text = preprocess_text(text, **kwargs)
|
|
tokenized = tokenizer(text, padding="max_length", truncation=True,
|
|
return_tensors="pt", max_length=max_length)
|
|
tokenized = {key: value.to(device) for key, value in tokenized.items()}
|
|
y = text_encoder(**tokenized).last_hidden_state
|
|
y = y * tokenized["attention_mask"].unsqueeze(-1)
|
|
return (y, tokenized) if return_tokens else y
|
|
|
|
|
|
@torch.no_grad()
|
|
def add_noise(x, noise_std=0.3, noise_type="gaussian", cond_noise_level=False):
|
|
if isinstance(x, list):
|
|
return zip(*[add_noise(xi, noise_std, noise_type) for xi in x])
|
|
if noise_type == "gaussian":
|
|
noise = noise_std * torch.randn_like(x)
|
|
return x + noise, noise
|
|
if noise_type == "uniform":
|
|
noise = torch.rand_like(x)
|
|
return ((x + 1) * (255 / 2) + noise) / 256 * 2 - 1, noise
|
|
raise NotImplementedError(noise_type)
|
|
|
|
|
|
# Import ml_starflow components
|
|
from .ml_starflow.utils.model_setup import setup_transformer
|
|
try:
|
|
from .starflow_config import get_config_args, get_config_names
|
|
except Exception:
|
|
_FALLBACK_CONFIGS = {
|
|
"starflow_3B_t2i_256x256": {
|
|
"no_flip": 1,
|
|
"fsdp": 1,
|
|
"fsdp_text_encoder": 1,
|
|
"img_size": 256,
|
|
"txt_size": 128,
|
|
"channel_size": 4,
|
|
"patch_size": 1,
|
|
"channels": 3072,
|
|
"blocks": 6,
|
|
"layers_per_block": [2, 2, 2, 2, 2, 24],
|
|
"noise_std": 0.3,
|
|
"batch_size": 1024,
|
|
"lr": 6.4e-05,
|
|
"min_lr": 1e-6,
|
|
"nvp": 1,
|
|
"rope": 1,
|
|
"adaln": 0,
|
|
"sos": 1,
|
|
"seq_order": "L2R",
|
|
"wds": 1,
|
|
"use_softplus": 1,
|
|
"cond_top_only": 1,
|
|
"use_final_norm": 1,
|
|
"gradient_checkpoint": 0,
|
|
"vae": "stabilityai/sd-vae-ft-ema",
|
|
"text": "google/flan-t5-xl",
|
|
"cfg": 2.5,
|
|
"drop_label": 0.1,
|
|
"sample_freq": 10,
|
|
"soft_clip": 4.0,
|
|
"latent_norm_regularization": 1e-4,
|
|
}
|
|
}
|
|
|
|
def get_config_names():
|
|
return sorted(_FALLBACK_CONFIGS.keys())
|
|
|
|
def get_config_args(name):
|
|
if name not in _FALLBACK_CONFIGS:
|
|
raise KeyError(f"Unknown STARFlow config preset: {name}")
|
|
return _FALLBACK_CONFIGS[name]
|
|
|
|
|
|
_DEFAULT_ARGS = {
|
|
"channel_size": 3,
|
|
"img_size": 32,
|
|
"txt_size": 0,
|
|
"sos": 0,
|
|
"cond_top_only": 0,
|
|
"use_softplus": 0,
|
|
"use_pretrained_lm": None,
|
|
"use_mm_attn": 0,
|
|
"use_final_norm": 0,
|
|
"soft_clip": 0.0,
|
|
"seq_order": "R2L",
|
|
"learnable_self_denoiser": 0,
|
|
"conditional_denoiser": 0,
|
|
"noise_embed_denoiser": 0,
|
|
"temporal_causal": 0,
|
|
"shallow_block_local": 0,
|
|
"denoiser_window": None,
|
|
"local_attn_window": None,
|
|
"top_block_channels": None,
|
|
"patch_size": 4,
|
|
"channels": 512,
|
|
"blocks": 4,
|
|
"layers_per_block": [8],
|
|
"rope": 0,
|
|
"pt_seq_len": None,
|
|
"head_dim": 64,
|
|
"num_heads": None,
|
|
"num_kv_heads": None,
|
|
"use_swiglu": 0,
|
|
"use_bias": 1,
|
|
"use_qk_norm": 0,
|
|
"use_post_norm": 0,
|
|
"norm_type": "layer_norm",
|
|
"vae": None,
|
|
"text": None,
|
|
}
|
|
|
|
_INT_KEYS = {
|
|
"channel_size",
|
|
"img_size",
|
|
"txt_size",
|
|
"sos",
|
|
"cond_top_only",
|
|
"use_softplus",
|
|
"use_mm_attn",
|
|
"use_final_norm",
|
|
"rope",
|
|
"patch_size",
|
|
"channels",
|
|
"blocks",
|
|
"head_dim",
|
|
"num_heads",
|
|
"num_kv_heads",
|
|
"use_swiglu",
|
|
"use_bias",
|
|
"use_qk_norm",
|
|
"use_post_norm",
|
|
"temporal_causal",
|
|
"learnable_self_denoiser",
|
|
"conditional_denoiser",
|
|
"noise_embed_denoiser",
|
|
"shallow_block_local",
|
|
"top_block_channels",
|
|
"pt_seq_len",
|
|
"denoiser_window",
|
|
"local_attn_window",
|
|
}
|
|
|
|
_FLOAT_KEYS = {
|
|
"soft_clip",
|
|
}
|
|
|
|
|
|
def _parse_layers_per_block(value):
|
|
if isinstance(value, str):
|
|
parts = [p for p in value.split() if p]
|
|
return [int(p) for p in parts]
|
|
if isinstance(value, (list, tuple)):
|
|
return [int(v) for v in value]
|
|
if isinstance(value, int):
|
|
return [value]
|
|
return value
|
|
|
|
|
|
def _parse_config_args(config_dict):
|
|
return dict(config_dict)
|
|
|
|
|
|
def _build_args_from_config(config_args):
|
|
args = dict(_DEFAULT_ARGS)
|
|
for key, value in config_args.items():
|
|
if key == "layers_per_block":
|
|
args[key] = _parse_layers_per_block(value)
|
|
continue
|
|
if value is None:
|
|
args[key] = None
|
|
continue
|
|
if key in _INT_KEYS:
|
|
args[key] = int(value)
|
|
continue
|
|
if key in _FLOAT_KEYS:
|
|
args[key] = float(value)
|
|
continue
|
|
args[key] = value
|
|
return SimpleNamespace(**args)
|
|
|
|
|
|
def _load_state_dict(path):
|
|
if path.endswith(".safetensors"):
|
|
return load_file(path)
|
|
return torch.load(path, map_location="cpu")
|
|
|
|
def _get_checkpoint_choices():
|
|
seen = set()
|
|
choices = []
|
|
for folder_name in ("starflow_checkpoints", "diffusion_models", "checkpoints"):
|
|
try:
|
|
names = folder_paths.get_filename_list(folder_name)
|
|
except Exception:
|
|
names = []
|
|
for name in names:
|
|
if name in seen:
|
|
continue
|
|
seen.add(name)
|
|
choices.append(name)
|
|
return sorted(choices)
|
|
|
|
|
|
def _resolve_checkpoint_path(checkpoint_name):
|
|
for folder_name in ("starflow_checkpoints", "diffusion_models", "checkpoints"):
|
|
path = folder_paths.get_full_path(folder_name, checkpoint_name)
|
|
if path:
|
|
return path
|
|
raise FileNotFoundError(f"Checkpoint '{checkpoint_name}' not found in STARFlow, diffusion_models, or checkpoints.")
|
|
|
|
|
|
def _get_vae_downscale_ratio(vae, fallback=8):
|
|
ratio = getattr(vae, "downscale_ratio", None)
|
|
if isinstance(ratio, int):
|
|
return ratio
|
|
if isinstance(ratio, tuple):
|
|
ints = [v for v in ratio if isinstance(v, int)]
|
|
if ints:
|
|
return int(ints[-1])
|
|
return fallback
|
|
|
|
|
|
def _prepare_input_image(image, height, width, vae, device, noise_std):
|
|
if image is None:
|
|
return None
|
|
img = image.to(device=device)
|
|
if img.dim() != 4:
|
|
raise ValueError("input_image must be a 4D tensor (B,H,W,C)")
|
|
b, h, w, c = img.shape
|
|
if c != 3:
|
|
img = img[..., :3]
|
|
img = img.permute(0, 3, 1, 2).contiguous()
|
|
scale = max(height / h, width / w)
|
|
new_h = max(1, int(round(h * scale)))
|
|
new_w = max(1, int(round(w * scale)))
|
|
if new_h != h or new_w != w:
|
|
img = F.interpolate(img, size=(new_h, new_w), mode="bilinear", align_corners=False)
|
|
top = max(0, (new_h - height) // 2)
|
|
left = max(0, (new_w - width) // 2)
|
|
img = img[:, :, top:top + height, left:left + width]
|
|
img = img.permute(0, 2, 3, 1).contiguous()
|
|
latent = vae.encode(img)
|
|
latent, _ = add_noise(latent, noise_std)
|
|
return latent
|
|
|
|
|
|
def _apply_denoising(model, x_chunk, y_batch, text_encoder, tokenizer, args,
|
|
text_encoder_kwargs, sigma_curr, sigma_next=0):
|
|
noise_std_const = 0.3 # constant used for noise levels in official code
|
|
if isinstance(y_batch, torch.Tensor):
|
|
y_ = y_batch
|
|
elif y_batch is not None:
|
|
y_ = encode_text(
|
|
text_encoder,
|
|
tokenizer,
|
|
y_batch,
|
|
args.txt_size,
|
|
text_encoder.device,
|
|
**text_encoder_kwargs,
|
|
)
|
|
else:
|
|
y_ = None
|
|
|
|
if getattr(args, "disable_learnable_denoiser", False) or not hasattr(
|
|
model, "learnable_self_denoiser"
|
|
):
|
|
return _self_denoise(
|
|
model,
|
|
x_chunk,
|
|
y_,
|
|
noise_std=sigma_curr,
|
|
steps=1,
|
|
disable_learnable_denoiser=getattr(
|
|
args, "disable_learnable_denoiser", False
|
|
),
|
|
)
|
|
|
|
if sigma_curr is not None and isinstance(y_batch, (list, type(None))):
|
|
text_encoder_kwargs["noise_std"] = sigma_curr
|
|
denoiser_output = model(x_chunk, y_, denoiser=True)
|
|
return x_chunk - denoiser_output * noise_std_const * (sigma_curr - sigma_next) / sigma_curr
|
|
|
|
|
|
def _self_denoise(model, samples, y, noise_std=0.1, lr=1, steps=1,
|
|
disable_learnable_denoiser=False):
|
|
if steps == 0:
|
|
return samples
|
|
|
|
outputs = []
|
|
x = samples.clone()
|
|
lr = noise_std ** 2 * lr
|
|
with torch.enable_grad():
|
|
x.requires_grad = True
|
|
model.train()
|
|
z, _, _, logdets = model(x, y)
|
|
loss = model.get_loss(z, logdets)["loss"] * 65536
|
|
grad = float(samples.numel()) / 65536 * torch.autograd.grad(loss, [x])[0]
|
|
outputs += [(x - grad * lr).detach()]
|
|
x = torch.cat(outputs, -1)
|
|
return x
|
|
|
|
|
|
def _process_denoising(samples, y, args, model, text_encoder, tokenizer,
|
|
text_encoder_kwargs, noise_std):
|
|
if not (
|
|
args.finetuned_vae == "none"
|
|
and getattr(args, "vae_adapter", None) is None
|
|
and getattr(args, "return_sequence", 0) == 0
|
|
):
|
|
return samples
|
|
|
|
torch.cuda.empty_cache()
|
|
assert isinstance(samples, torch.Tensor)
|
|
samples = samples.cpu()
|
|
|
|
b = samples.size(0)
|
|
db = max(1, min(getattr(args, "denoising_batch_size", 1), b))
|
|
denoised_samples = []
|
|
is_video = samples.dim() == 5
|
|
|
|
for j in range(b // db):
|
|
x_all = torch.clone(samples[j * db : (j + 1) * db]).detach().cuda()
|
|
y_batch = y[j * db : (j + 1) * db] if y is not None else None
|
|
|
|
if is_video:
|
|
s_idx, overlap = 0, 0
|
|
steps = (
|
|
x_all.size(1)
|
|
if getattr(args, "local_attn_window", None) is None
|
|
else args.local_attn_window
|
|
)
|
|
|
|
while s_idx < x_all.size(1):
|
|
x_chunk = x_all[:, s_idx : s_idx + steps].detach().clone()
|
|
x_denoised = _apply_denoising(
|
|
model,
|
|
x_chunk,
|
|
y_batch,
|
|
text_encoder,
|
|
tokenizer,
|
|
args,
|
|
text_encoder_kwargs,
|
|
noise_std,
|
|
)
|
|
x_all[:, s_idx + overlap : s_idx + steps] = x_denoised[:, overlap:]
|
|
overlap = (
|
|
steps - 1
|
|
if getattr(args, "denoiser_window", None) is None
|
|
else args.denoiser_window
|
|
)
|
|
s_idx += steps - overlap
|
|
else:
|
|
x_all = _apply_denoising(
|
|
model,
|
|
x_all,
|
|
y_batch,
|
|
text_encoder,
|
|
tokenizer,
|
|
args,
|
|
text_encoder_kwargs,
|
|
noise_std,
|
|
)
|
|
|
|
torch.cuda.empty_cache()
|
|
denoised_samples.append(x_all.detach().cpu())
|
|
|
|
return torch.cat(denoised_samples, dim=0).cuda()
|
|
|
|
|
|
class STARFlowCheckpointLoader:
|
|
"""Loads STARFlow transformer checkpoint and configuration"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Add starflow checkpoint folder
|
|
if not hasattr(folder_paths, '_starflow_folder_added'):
|
|
folder_paths.add_model_folder_path("starflow_checkpoints", os.path.join(folder_paths.models_dir, "starflow"))
|
|
for path in folder_paths.get_folder_paths("diffusion_models"):
|
|
folder_paths.add_model_folder_path("starflow_checkpoints", path)
|
|
folder_paths._starflow_folder_added = True
|
|
|
|
return {
|
|
"required": {
|
|
"checkpoint_name": (_get_checkpoint_choices(),
|
|
{"tooltip": "STARFlow checkpoint file (.pth or .safetensors)"}),
|
|
"config_name": (get_config_names(),
|
|
{"tooltip": "STARFlow architecture preset"}),
|
|
},
|
|
"optional": {
|
|
"vae_downscale": ("INT", {
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 64,
|
|
"tooltip": "VAE downscale ratio used during training (8 for SD VAEs)"
|
|
}),
|
|
"text_dim": ("INT", {
|
|
"default": 2048,
|
|
"min": 1,
|
|
"max": 8192,
|
|
"tooltip": "Text encoder hidden size (2048 for FLAN-T5-XL)"
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STARFLOW_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "STARFlow"
|
|
DESCRIPTION = "Loads a STARFlow transformer model checkpoint with configuration"
|
|
|
|
def load_checkpoint(self, checkpoint_name, config_name, vae_downscale=8, text_dim=2048):
|
|
# Load config preset
|
|
config_args = _parse_config_args(get_config_args(config_name))
|
|
args = _build_args_from_config(config_args)
|
|
|
|
# Initialize model
|
|
device = mm.get_torch_device()
|
|
|
|
if args.img_size % vae_downscale != 0:
|
|
raise ValueError(
|
|
f"img_size {args.img_size} must be divisible by vae_downscale {vae_downscale}"
|
|
)
|
|
latent_img_size = args.img_size // vae_downscale
|
|
model_args = SimpleNamespace(**{**vars(args), "img_size": latent_img_size})
|
|
model = setup_transformer(
|
|
model_args,
|
|
dist=None,
|
|
txt_dim=text_dim,
|
|
use_checkpoint=0,
|
|
use_checkpoint_mlp=None,
|
|
)
|
|
|
|
# Load checkpoint
|
|
checkpoint_path = _resolve_checkpoint_path(checkpoint_name)
|
|
state_dict = _load_state_dict(checkpoint_path)
|
|
|
|
# Handle different checkpoint formats
|
|
if 'model' in state_dict:
|
|
state_dict = state_dict['model']
|
|
if 'state_dict' in state_dict:
|
|
state_dict = state_dict['state_dict']
|
|
|
|
model.load_state_dict(state_dict, strict=False)
|
|
model.eval()
|
|
model = model.to(device)
|
|
|
|
# Disable gradients
|
|
for param in model.parameters():
|
|
param.requires_grad = False
|
|
|
|
return ({
|
|
"model": model,
|
|
"config": args,
|
|
"image_size": args.img_size,
|
|
"latent_size": latent_img_size,
|
|
"channel_size": args.channel_size,
|
|
"txt_size": args.txt_size,
|
|
"text_dim": text_dim,
|
|
"vae_downscale": vae_downscale,
|
|
},)
|
|
|
|
|
|
class STARFlowT5TextEncode:
|
|
"""Encodes text using T5-XL (FLAN-T5-XL) and prepares STARFlow conditioning"""
|
|
|
|
# Class-level cache for T5 model
|
|
_t5_cache = {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"t5_model": (folder_paths.get_filename_list("text_encoders"),
|
|
{"tooltip": "T5-XL model file from models/text_encoders/ (e.g., flan-t5-xl.safetensors)"}),
|
|
"text": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"forceInput": True,
|
|
"tooltip": "Text prompt to encode (connect a String node)"
|
|
}),
|
|
"aspect_ratio": (["1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21"],
|
|
{"default": "1:1", "tooltip": "Output aspect ratio"}),
|
|
},
|
|
"optional": {
|
|
"max_length": ("INT", {
|
|
"default": 128,
|
|
"min": 64,
|
|
"max": 512,
|
|
"tooltip": "Maximum sequence length for T5 tokenizer"
|
|
}),
|
|
"noise_std": ("FLOAT", {
|
|
"default": 0.3,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"tooltip": "Noise standard deviation for generation"
|
|
}),
|
|
"cond_noise_level": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Include noise level in text conditioning"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STARFLOW_CONDITIONING",)
|
|
RETURN_NAMES = ("conditioning",)
|
|
FUNCTION = "encode_text"
|
|
CATEGORY = "STARFlow"
|
|
DESCRIPTION = "Encodes text using T5-XL (google/flan-t5-xl) encoder with STARFlow-specific metadata"
|
|
|
|
def encode_text(self, t5_model, text, aspect_ratio, max_length=128, noise_std=0.3, cond_noise_level=False):
|
|
device = mm.text_encoder_device()
|
|
dtype = mm.text_encoder_dtype(device)
|
|
|
|
resolved_path = folder_paths.get_full_path_or_raise("text_encoders", t5_model)
|
|
t5_repo = "google/flan-t5-xl"
|
|
|
|
cache_key = resolved_path
|
|
# Load T5-XL model (with caching)
|
|
if cache_key not in self._t5_cache:
|
|
print(f"Loading T5-XL model: {resolved_path}")
|
|
|
|
# Initialize T5-XL model architecture (not XXL!)
|
|
# T5-XL has: d_model=2048, num_heads=32, d_ff=5120
|
|
config = AutoConfig.from_pretrained(t5_repo)
|
|
text_encoder = T5EncoderModel(config)
|
|
|
|
# Load weights from safetensors
|
|
if resolved_path.endswith('.safetensors'):
|
|
state_dict = load_file(resolved_path)
|
|
if any(k.startswith("decoder.") for k in state_dict.keys()):
|
|
print(
|
|
"Warning: T5 weights look like decoder weights. "
|
|
"Use an encoder-only file (e.g., models--google--flan-t5-xl--text_encoder.safetensors)."
|
|
)
|
|
missing, unexpected = text_encoder.load_state_dict(state_dict, strict=False)
|
|
if missing or unexpected:
|
|
print(f"T5 load_state_dict missing={len(missing)} unexpected={len(unexpected)}")
|
|
else:
|
|
# Try loading as pytorch checkpoint
|
|
state_dict = torch.load(resolved_path, map_location="cpu")
|
|
if 'model' in state_dict:
|
|
state_dict = state_dict['model']
|
|
text_encoder.load_state_dict(state_dict, strict=False)
|
|
|
|
# Load tokenizer (uses HF cache, very small download)
|
|
tokenizer = T5Tokenizer.from_pretrained(t5_repo)
|
|
|
|
text_encoder.to(device=device, dtype=dtype)
|
|
text_encoder.eval()
|
|
|
|
# Disable gradients
|
|
for param in text_encoder.parameters():
|
|
param.requires_grad = False
|
|
|
|
self._t5_cache[cache_key] = (tokenizer, text_encoder)
|
|
print("T5-XL model loaded successfully!")
|
|
|
|
tokenizer, text_encoder = self._t5_cache[cache_key]
|
|
|
|
# Preprocess text with aspect ratio (STARFlow style)
|
|
text_list = [text]
|
|
processed_text = preprocess_text(
|
|
text_list,
|
|
aspect_ratio=aspect_ratio,
|
|
noise_std=noise_std if cond_noise_level else None,
|
|
)
|
|
|
|
# Tokenize
|
|
tokenized = tokenizer(
|
|
processed_text,
|
|
padding="max_length",
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
max_length=max_length
|
|
)
|
|
tokenized = {k: v.to(device) for k, v in tokenized.items()}
|
|
|
|
# Encode text
|
|
with torch.no_grad():
|
|
# Positive embedding
|
|
embeddings_pos = text_encoder(**tokenized).last_hidden_state
|
|
# Mask padding tokens
|
|
embeddings_pos = embeddings_pos * tokenized['attention_mask'].unsqueeze(-1)
|
|
|
|
# Negative embedding (empty prompt for CFG)
|
|
empty_text = preprocess_text(
|
|
[""],
|
|
aspect_ratio=aspect_ratio,
|
|
noise_std=noise_std if cond_noise_level else None,
|
|
)
|
|
empty_tokenized = tokenizer(
|
|
empty_text,
|
|
padding="max_length",
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
max_length=max_length
|
|
)
|
|
empty_tokenized = {k: v.to(device) for k, v in empty_tokenized.items()}
|
|
embeddings_neg = text_encoder(**empty_tokenized).last_hidden_state
|
|
embeddings_neg = embeddings_neg * empty_tokenized['attention_mask'].unsqueeze(-1)
|
|
|
|
# Concatenate for CFG: [positive, negative]
|
|
embeddings = torch.cat([embeddings_pos, embeddings_neg], dim=0)
|
|
|
|
# Create conditioning dict with metadata
|
|
conditioning = {
|
|
"embeddings": embeddings,
|
|
"metadata": {
|
|
"aspect_ratio": aspect_ratio,
|
|
"noise_std": noise_std,
|
|
"cond_noise_level": cond_noise_level,
|
|
"text": text,
|
|
}
|
|
}
|
|
conditioning["tokenizer"] = tokenizer
|
|
conditioning["text_encoder"] = text_encoder
|
|
|
|
return (conditioning,)
|
|
|
|
|
|
class STARFlowSampler:
|
|
"""Generates latent samples using STARFlow model"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("STARFLOW_MODEL", {"tooltip": "STARFlow model"}),
|
|
"conditioning": ("STARFLOW_CONDITIONING",
|
|
{"tooltip": "Text conditioning from T5 encoder"}),
|
|
"vae": ("VAE", {"tooltip": "VAE for determining latent dimensions"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff,
|
|
"tooltip": "Random seed for reproducible generation"}),
|
|
"cfg": ("FLOAT", {
|
|
"default": 2.5,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"tooltip": "Classifier-free guidance scale"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"input_image": ("IMAGE", {
|
|
"tooltip": "Optional image for conditioning (uses VAE encode)"
|
|
}),
|
|
"jacobi": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Enable Jacobi iteration for faster sampling"
|
|
}),
|
|
"jacobi_th": ("FLOAT", {
|
|
"default": 0.001,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.0001,
|
|
"tooltip": "Jacobi convergence threshold (0.001 recommended for images)"
|
|
}),
|
|
"jacobi_block_size": ("INT", {
|
|
"default": 16,
|
|
"min": 1,
|
|
"max": 128,
|
|
"tooltip": "Jacobi block size (16 recommended for images)"
|
|
}),
|
|
"jacobi_max_iter": ("INT", {
|
|
"default": 32,
|
|
"min": 1,
|
|
"max": 100,
|
|
"tooltip": "Maximum Jacobi iterations"
|
|
}),
|
|
"attn_temp": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": -1.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"tooltip": "Attention temperature (for non-Jacobi mode)"
|
|
}),
|
|
"guide_top": ("INT", {
|
|
"default": -1,
|
|
"min": -1,
|
|
"max": 1024,
|
|
"tooltip": "Guide top blocks only (-1 disables)"
|
|
}),
|
|
"annealed_guidance": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Anneal guidance during sampling"
|
|
}),
|
|
"return_sequence": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Return the full reverse sequence (debug)"
|
|
}),
|
|
"denoising_batch_size": ("INT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 128,
|
|
"tooltip": "Denoising batch size (official behavior)"
|
|
}),
|
|
"disable_learnable_denoiser": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Disable learnable denoiser if present"
|
|
}),
|
|
"target_noise_std": ("FLOAT", {
|
|
"default": -1.0,
|
|
"min": -1000.0,
|
|
"max": 1000.0,
|
|
"step": 0.01,
|
|
"tooltip": "Override noise_std when >= 0"
|
|
}),
|
|
"context_length": ("INT", {
|
|
"default": 16,
|
|
"min": 0,
|
|
"max": 1024,
|
|
"tooltip": "Context length for video Jacobi sampling"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("samples",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "STARFlow"
|
|
DESCRIPTION = "Generates latent samples using STARFlow reverse flow"
|
|
|
|
def sample(self, model, conditioning, vae, seed, cfg,
|
|
input_image=None, jacobi=False, jacobi_th=0.001, jacobi_block_size=16,
|
|
jacobi_max_iter=32, attn_temp=1.0, guide_top=-1,
|
|
annealed_guidance=False,
|
|
return_sequence=False, denoising_batch_size=1,
|
|
disable_learnable_denoiser=False, target_noise_std=-1.0,
|
|
context_length=16):
|
|
|
|
# Set random seed
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
np.random.seed(seed)
|
|
|
|
# Extract model and config
|
|
starflow_model = model["model"]
|
|
config = model["config"]
|
|
|
|
# Get metadata
|
|
metadata = conditioning["metadata"]
|
|
aspect_ratio = metadata["aspect_ratio"]
|
|
noise_std = metadata.get("noise_std", 0.3)
|
|
if target_noise_std is not None and float(target_noise_std) >= 0:
|
|
noise_std = float(target_noise_std)
|
|
|
|
# Calculate image dimensions based on aspect ratio
|
|
base_size = model.get("image_size", getattr(config, "img_size", 256))
|
|
vae_downscale = _get_vae_downscale_ratio(vae, model.get("vae_downscale", 8))
|
|
height, width = aspect_ratio_to_image_size(base_size, aspect_ratio, multiple=vae_downscale)
|
|
|
|
# Calculate latent dimensions (VAE downsamples by 8x)
|
|
latent_h = height // vae_downscale
|
|
latent_w = width // vae_downscale
|
|
|
|
# Get embeddings
|
|
embeddings = conditioning["embeddings"]
|
|
if cfg <= 0:
|
|
embeddings = embeddings[:1]
|
|
|
|
# Generate Gaussian noise
|
|
device = mm.get_torch_device()
|
|
latent_channels = getattr(vae, "latent_channels", model.get("channel_size", 4))
|
|
model_dtype = next(starflow_model.parameters()).dtype
|
|
batch = 1 if cfg > 0 else embeddings.shape[0]
|
|
noise = torch.randn(
|
|
batch,
|
|
latent_channels,
|
|
latent_h,
|
|
latent_w,
|
|
device=device,
|
|
dtype=model_dtype,
|
|
)
|
|
|
|
# Prepare sampling kwargs
|
|
sampling_kwargs = {
|
|
"reverse": True,
|
|
"guidance": cfg,
|
|
"jacobi": 1 if jacobi else 0,
|
|
"context_length": context_length,
|
|
"return_sequence": bool(return_sequence),
|
|
}
|
|
if guide_top is not None and int(guide_top) >= 0:
|
|
sampling_kwargs["guide_top"] = int(guide_top)
|
|
|
|
if jacobi:
|
|
if jacobi_th > 0.1:
|
|
print(f"Warning: jacobi_th {jacobi_th} is unusually high; clamping to 0.1")
|
|
jacobi_th = 0.1
|
|
sampling_kwargs.update({
|
|
"jacobi": 1,
|
|
"jacobi_th": jacobi_th,
|
|
"jacobi_block_size": jacobi_block_size,
|
|
"jacobi_max_iter": jacobi_max_iter,
|
|
})
|
|
else:
|
|
sampling_kwargs.update({
|
|
"attn_temp": attn_temp,
|
|
"annealed_guidance": bool(annealed_guidance),
|
|
})
|
|
sampling_kwargs["verbose"] = True
|
|
|
|
embeddings = embeddings.to(device=device, dtype=model_dtype)
|
|
|
|
kv_caches = None
|
|
if input_image is not None:
|
|
if cfg <= 0:
|
|
raise ValueError("CFG must be > 0 for image-conditioned sampling")
|
|
input_latent = _prepare_input_image(
|
|
input_image, height, width, vae, device, noise_std
|
|
)
|
|
input_latent = input_latent.to(device=device, dtype=model_dtype)
|
|
kv_caches = starflow_model(input_latent.unsqueeze(1), embeddings, context=True)
|
|
|
|
# Forward pass (reverse mode for generation)
|
|
with torch.no_grad():
|
|
samples = starflow_model(noise, embeddings, kv_caches=kv_caches, **sampling_kwargs)
|
|
if kv_caches is not None:
|
|
del kv_caches
|
|
torch.cuda.empty_cache()
|
|
|
|
tokenizer = conditioning.get("tokenizer")
|
|
text_encoder = conditioning.get("text_encoder")
|
|
if isinstance(samples, torch.Tensor) and tokenizer is not None and text_encoder is not None:
|
|
cond_noise_level = bool(metadata.get("cond_noise_level", False))
|
|
text_encoder_kwargs = {
|
|
"aspect_ratio": aspect_ratio,
|
|
"noise_std": noise_std if cond_noise_level else None,
|
|
}
|
|
args = SimpleNamespace(
|
|
finetuned_vae="",
|
|
vae_adapter=None,
|
|
return_sequence=1 if return_sequence else 0,
|
|
denoising_batch_size=int(denoising_batch_size),
|
|
disable_learnable_denoiser=bool(disable_learnable_denoiser),
|
|
denoiser_window=getattr(config, "denoiser_window", None),
|
|
local_attn_window=getattr(config, "local_attn_window", None),
|
|
txt_size=model.get("txt_size", getattr(config, "txt_size", 128)),
|
|
)
|
|
y_caption = [metadata.get("text", "")] * samples.size(0)
|
|
samples = _process_denoising(
|
|
samples,
|
|
y_caption,
|
|
args,
|
|
starflow_model,
|
|
text_encoder,
|
|
tokenizer,
|
|
text_encoder_kwargs,
|
|
noise_std,
|
|
)
|
|
|
|
# Return in ComfyUI LATENT format
|
|
return ({"samples": samples},)
|
|
|
|
|
|
class STARFlowVAEDecode:
|
|
"""Decodes STARFlow latents to images using ComfyUI VAE"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT", {"tooltip": "Latent samples to decode"}),
|
|
"vae": ("VAE", {"tooltip": "VAE decoder (use ComfyUI's VAELoader)"}),
|
|
},
|
|
"optional": {
|
|
"vae_scaling": ("FLOAT", {
|
|
"default": 0.18215,
|
|
"min": 0.01,
|
|
"max": 2.0,
|
|
"step": 0.00001,
|
|
"tooltip": "VAE scaling factor (0.18215 for SD1.x/SD2.x, 0.13025 for SDXL)"
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "decode"
|
|
CATEGORY = "STARFlow"
|
|
DESCRIPTION = "Decodes STARFlow latents to images"
|
|
|
|
def decode(self, samples, vae, vae_scaling=0.18215):
|
|
# Extract latent tensor
|
|
latent = samples["samples"]
|
|
|
|
# Decode using VAE
|
|
with torch.no_grad():
|
|
scaled = latent / float(vae_scaling) if vae_scaling else latent
|
|
image = vae.decode(scaled)
|
|
|
|
return (image,)
|