""" 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 = "user\nPlease generate {mode} about: {prompt}\nmodel\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,)