From df93865bf5a2e7854b8c28dd3f4e3f3963baa6d9 Mon Sep 17 00:00:00 2001 From: MaximClouser Date: Wed, 26 Feb 2025 19:05:09 +0000 Subject: [PATCH] zero order search, loading verifiers --- src/inferencescale/nodes.py | 678 +++++++++++++++++++++++++++--------- 1 file changed, 513 insertions(+), 165 deletions(-) diff --git a/src/inferencescale/nodes.py b/src/inferencescale/nodes.py index 757d1d3..7ccb9de 100644 --- a/src/inferencescale/nodes.py +++ b/src/inferencescale/nodes.py @@ -1,224 +1,572 @@ +import json +import logging +import os +import traceback +from typing import Dict, List, Optional, Tuple, Any + +import ImageReward as reward import torch import torch.nn.functional as F -import comfy.samplers +from huggingface_hub import snapshot_download +from torchvision.transforms import ToPILImage +from torchvision.utils import make_grid +from transformers import CLIPProcessor, CLIPModel + import comfy.sample +import comfy.samplers import comfy.sd import comfy.utils +import folder_paths import latent_preview -import json -from torchvision.utils import make_grid -from torchvision.transforms import ToPILImage +# from comfy.utils import ProgressBar -from transformers import CLIPProcessor, CLIPModel -import ImageReward as reward +from .qwen_verifier import QwenVLMVerifier +from .utils import generate_neighbors + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger('inferencescale') -def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False): +def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, + disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, + override_noise=None): latent_image = latent["samples"] latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image) - - if disable_noise: + if override_noise is not None: + noise = override_noise.to(latent_image.device) + elif disable_noise: noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") else: - batch_inds = latent["batch_index"] if "batch_index" in latent else None + batch_inds = latent.get("batch_index", None) noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds) - noise_mask = None - if "noise_mask" in latent: - noise_mask = latent["noise_mask"] + noise_mask = latent.get("noise_mask", None) + # callback = latent_preview.prepare_callback(model, steps) + # disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + disable_pbar = True + callback = None + + samples = comfy.sample.sample( + model=model, + noise=noise, + steps=steps, + cfg=cfg, + sampler_name=sampler_name, + scheduler=scheduler, + positive=positive, + negative=negative, + latent_image=latent_image, + denoise=denoise, + disable_noise=disable_noise, + start_step=start_step, + last_step=last_step, + force_full_denoise=force_full_denoise, + noise_mask=noise_mask, + callback=callback, + disable_pbar=disable_pbar, + seed=seed + ) - callback = latent_preview.prepare_callback(model, steps) - disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED - samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, - denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step, - force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) out = latent.copy() out["samples"] = samples return (out, ) -def compute_similarity(clip_model, clip_processor, text, image) -> float: - # Process inputs: 'text' is a string, and 'image' should be a PIL image - inputs = clip_processor(text=[text], images=image, return_tensors="pt", padding=True) - - device = next(clip_model.parameters()).device - inputs = {k: v.to(device) for k, v in inputs.items()} - - # Get text and image embeddings. - with torch.no_grad(): - text_embeddings = clip_model.get_text_features( - input_ids=inputs['input_ids'], - attention_mask=inputs['attention_mask'] - ) - image_embeddings = clip_model.get_image_features( - pixel_values=inputs['pixel_values'] - ) - - text_embeddings = F.normalize(text_embeddings, dim=-1) - image_embeddings = F.normalize(image_embeddings, dim=-1) - - similarity = F.cosine_similarity(text_embeddings, image_embeddings) - # rounded_similarity = round(similarity.item(), 4) - rounded_similarity = similarity.item() - return rounded_similarity - - -def compute_combined_verifier_score(candidate_image, text_prompt, verifiers, weights): +def rank_candidates(results_by_verifier, num_candidates): """ - verifiers: a dict with keys like 'clip' and 'image_reward' and values holding the corresponding model(s)/processor(s) - weights: a dict with the same keys, providing a float weight for each verifier - Returns a combined score and a dict of individual scores. + results_by_verifier: { + 'clip': [score0, score1, ...], + 'image_reward': [...], + 'qwen_vlm_verifier': [...], + ... + } + Convert each verifier's raw scores into ranks (best=rank 1), then compute + an average rank for each candidate, returning a list of dicts sorted by avg_rank. """ - scores = {} + ranks_by_verifier = {} + for verifier_name, scores_list in results_by_verifier.items(): + # Sort indices by descending score (best score = rank 1) + sorted_indices = sorted(range(len(scores_list)), key=lambda i: scores_list[i], reverse=True) + rank_array = [0] * len(scores_list) + for rank, idx in enumerate(sorted_indices): + rank_array[idx] = rank + 1 # best = rank 1 + ranks_by_verifier[verifier_name] = rank_array - if 'clip' in verifiers and verifiers['clip'] is not None: - clip_model, clip_processor = verifiers['clip'] - clip_score = compute_similarity(clip_model, clip_processor, text_prompt, candidate_image) - scores['clip'] = clip_score - - if 'image_reward' in verifiers and verifiers['image_reward'] is not None: - image_reward_model = verifiers['image_reward'] - image_reward_score = image_reward_model.score(text_prompt, candidate_image) - scores['image_reward'] = image_reward_score - - combined_score = 0.0 - for key, score in scores.items(): - # TODO Implement algo from paper to combine scores - combined_score += weights.get(key, 1.0) * score - combined_score = round(combined_score, 4) - return combined_score, scores + n_verifiers = len(results_by_verifier) + final_results = [] + for i in range(num_candidates): + sum_ranks = 0 + image_scores = {} + for verifier_name, score_list in results_by_verifier.items(): + image_scores[verifier_name] = score_list[i] + sum_ranks += ranks_by_verifier[verifier_name][i] + avg_rank = (sum_ranks / n_verifiers) if n_verifiers > 0 else 9999 + final_results.append({"index": i, "avg_rank": avg_rank, "scores": image_scores}) + # Sort ascending by avg_rank (lower is better) + final_results.sort(key=lambda x: x["avg_rank"]) + return final_results -class MaximKSampler: +def score_candidates(candidate_tensors: List[torch.Tensor], text_prompt: str, verifier_names: Dict[str, Any]) -> List[Dict[str, float]]: + """ + Scores each candidate image using the provided verifier ensemble. + + Args: + candidate_tensors: List of image tensors to score + text_prompt: Text prompt to compare against + verifier_names: Dictionary of verifier instances + + Returns: + List of score dictionaries, one per candidate + """ + num_candidates = len(candidate_tensors) + raw_scores_list = [{} for _ in range(num_candidates)] + + try: + if not candidate_tensors: + raise ValueError("No candidate tensors provided for scoring") + + if not text_prompt: + logger.warning("Empty text prompt provided for scoring") + + for verifier_name, verifier in verifier_names.items(): + if verifier is None: + continue + + logger.info(f"Scoring {num_candidates} candidates with {verifier_name}") + + try: + with torch.no_grad(): + for i, tensor_img in enumerate(candidate_tensors): + try: + pil_img = ToPILImage()(tensor_img.squeeze(0).permute(2, 0, 1)) + + if verifier_name == "clip": + score = verifier.score(text_prompt, pil_img) + elif verifier_name == "image_reward": + score = verifier.score(text_prompt, pil_img) + elif verifier_name == "qwen_vlm_verifier": + score = verifier.get_overall_score(pil_img, text_prompt) + else: + logger.warning(f"Unknown verifier type: {verifier_name}") + continue + + raw_scores_list[i][verifier_name] = score + + except Exception as e: + logger.error(f"Error scoring candidate {i} with {verifier_name}: {str(e)}") + logger.debug(traceback.format_exc()) + raw_scores_list[i][verifier_name] = float('-inf') + + except Exception as e: + logger.error(f"Error in verifier {verifier_name}: {str(e)}") + logger.debug(traceback.format_exc()) + + except Exception as e: + logger.error(f"Fatal error in score_candidates: {str(e)}") + logger.debug(traceback.format_exc()) + raise + + return raw_scores_list + + +def rank_from_raw_scores(raw_scores_list): + """ + Converts a list of raw score dictionaries into a ranked list + using the existing ranking mechanism. + """ + # Build a verifier-to-scores dictionary from the raw_scores_list + results_by_verifier = {} + for candidate_scores in raw_scores_list: + for verifier_name, score in candidate_scores.items(): + if verifier_name not in results_by_verifier: + results_by_verifier[verifier_name] = [] + results_by_verifier[verifier_name].append(score) + num_candidates = len(raw_scores_list) + return rank_candidates(results_by_verifier, num_candidates) + + +class InferenceTimeScaler: + @classmethod def INPUT_TYPES(s): + return { "required": { - "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}), - "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}), - "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt; however, too high values may negatively impact quality."}), - "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The sampling algorithm used. This can affect the quality, speed, and style of the generated output."}), - "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}), - "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes to include in the image."}), - "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes to exclude from the image."}), - "latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}), - "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied; lower values maintain more of the structure of the initial image."}), - "text_prompt_to_compare": ("STRING", {"default": "", "multiline": True, "tooltip": "A text prompt used for comparing against generated images."}), - "search_budget": ("INT", {"default": 5, "min": 1, "max": 10000, "tooltip": "The number of random noise samples to generate (i.e., the search budget)."}), - "vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent image."}), - "top_k": ("INT", {"default": 3, "min": 1, "max": 100, "tooltip": "The top k image/score pairs to return."}), + "model": ("MODEL", {"tooltip": "Denoising model."}), + "vae": ("VAE", {"tooltip": "VAE model for decoding latents."}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "tooltip": "Random seed."}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "Number of denoising steps to apply during each forward evaluation."}), + "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "Classifier-Free Guidance Scale."}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "Sampling algorithm to be used during each forward evaluation."}), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "Noise removal scheduler."}), + "positive": ("CONDITIONING", {"tooltip": "Positive prompt conditioning."}), + "negative": ("CONDITIONING", {"tooltip": "Negative prompt conditioning."}), + "latent_image": ("LATENT", {"tooltip": "Latent image to be denoised."}), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Amount of denoising."}), + "search_algorithm": (["random", "zero-order"], {"default": "random", "tooltip": "Select the search algorithm: 'random' for standard random search or 'zero-order' for gradient-free local search."}), + "text_prompt_to_compare": ("STRING", {"default": "", "multiline": True, "tooltip": "Text prompt for verifier(s)."}), + "search_rounds": ("INT", {"default": 5, "min": 1, "max": 10000, "tooltip": "Number of search rounds (random seeds for random search, or iterations for zero-order search)."}), + "view_top_k": ("INT", {"default": 3, "min": 1, "max": 100, "tooltip": "Return grid view of the top-k images."}), + "num_neighbors": ("INT", {"default": 4, "min": 1, "max": 100, "tooltip": "Number of neighbors to sample per iteration in zero-order search (only used if search_algorithm is 'zero-order')."}), + "lambda_threshold": ("FLOAT", {"default": 0.90, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.001, "tooltip": "Perturbation step size for zero-order search (only used if search_algorithm is 'zero-order')."}) }, "optional": { - "clip_score_verifier": ("STRING", {"default": "openai/clip-vit-base-patch32", "tooltip": "Optional Hugging Face CLIP model identifier used for verifying the generated image quality."}), - "clip_score_weight": ("FLOAT", {"default": 1, "min": 0, "max": 1, "tooltip": "The weight of the clip verifier to apply "}), - "image_reward_verifier": ("STRING", {"default": "ImageReward-v1.0", "tooltip": "Optional Image Reward model used for verifying the generated image quality."}), - "image_reward_weight": ("FLOAT", {"default": 1, "min": 0, "max": 1, "tooltip": "The weight of the iamge rewqrd to apply "}) + "loaded_clip_score_verifier": ("CS_VERIFIER", {"tooltip": "HF CLIP model identifier."}), + "loaded_image_reward_verifier": ("IR_VERIFIER", {"tooltip": "ImageReward model identifier."}), + "loaded_qwen_verifier": ("QWN_VERIFIER", {"tooltip": "Qwen model identifier."}) } } RETURN_TYPES = ("IMAGE", "IMAGE", "STRING") RETURN_NAMES = ("Best Image", "Top-k Grid", "Top-k Score(s)") - - OUTPUT_TOOLTIPS = ("A grid of the top k images from the random search, and a JSON string with their verifier scores.",) + OUTPUT_TOOLTIPS = ("Best single image, grid of the top-k images, and their scores in JSON.",) FUNCTION = "execute" + CATEGORY = "InferenceTimeScaling" + DESCRIPTION = ( + "Performs inference-time optimization to find the best image matching your text prompt. " + "Supports two search algorithms:\n" + "1. Random Search: Generates multiple images with different random seeds\n" + "2. Zero-Order Search: Performs gradient-free local optimization\n\n" + "Uses an ensemble of AI verifiers (CLIP, ImageReward, Qwen-VL) to score and rank the generated images. " + "Returns the best image, a grid of top-k results, and detailed scoring information." + ) - CATEGORY = "maxtest/sampling" - DESCRIPTION = ("Uses the provided model, conditioning, and VAE to denoise the latent image. " - "Performs random search over noise candidates by sampling different noises, decodes each candidate image, " - "evaluates its quality using a verifier (if provided), and returns the top k candidates based on the verifier score.") + def execute(self, model, vae, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise, + text_prompt_to_compare, search_rounds, view_top_k, loaded_clip_score_verifier=None, loaded_image_reward_verifier=None, + search_algorithm="random", num_neighbors=4, lambda_threshold=0.01, loaded_qwen_verifier=None): + """ + Main execution method for inference time scaling. + """ + try: + # Only handling batch size of 1 for now + if latent_image["samples"].size(0) != 1: + raise ValueError(f"Expected latent image batch size of 1, got {latent_image['samples'].size(0)}") + logger.info(f"Starting InferenceTimeScaler with algorithm: {search_algorithm}") + logger.info(f"Search rounds: {search_rounds}, Top-k: {view_top_k}") - def execute(self, - model, - vae, - seed, - steps, - cfg, - sampler_name, - scheduler, - positive, - negative, - latent_image, - denoise, - text_prompt_to_compare, - search_budget, - top_k, - clip_score_verifier=None, - image_reward_verifier=None, - clip_score_weight=1.0, - image_reward_weight=1.0 - ): - # TODO Move to innit - device = "cuda" if torch.cuda.is_available() else "cpu" + if not text_prompt_to_compare: + logger.warning("Empty text prompt provided") - candidate_results = [] + use_zero_order = (search_algorithm == "zero-order") + + # Validate inputs + if steps < 1: + raise ValueError("Steps must be >= 1") + if cfg < 0: + raise ValueError("CFG must be >= 0") + if denoise < 0 or denoise > 1: + raise ValueError("Denoise must be between 0 and 1") - verifiers = {} - weights = {} - - if clip_score_verifier: - clip_model = CLIPModel.from_pretrained(clip_score_verifier) - clip_processor = CLIPProcessor.from_pretrained(clip_score_verifier) - clip_model.to(device) - verifiers['clip'] = (clip_model, clip_processor) - weights['clip'] = clip_score_weight + # Store loaded verifier models + verifier_names = { + "clip": loaded_clip_score_verifier, + "image_reward": loaded_image_reward_verifier, + "qwen_vlm_verifier": loaded_qwen_verifier + } - if image_reward_verifier: - image_reward_model = reward.load(image_reward_verifier) - verifiers['image_reward'] = image_reward_model - weights['image_reward'] = image_reward_weight + active_verifiers = sum(1 for v in verifier_names.values() if v is not None) + if active_verifiers == 0: + raise ValueError("No verifiers provided - at least one verifier is required") + + logger.info(f"Using {active_verifiers} active verifiers") - for i in range(search_budget): + if use_zero_order: + result = self._execute_zero_order( + model, vae, seed, steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, denoise, text_prompt_to_compare, + verifier_names, search_rounds, num_neighbors, lambda_threshold, view_top_k + ) + else: + result = self._execute_random_search( + model, vae, seed, steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, denoise, text_prompt_to_compare, + verifier_names, search_rounds, view_top_k + ) + + logger.info("InferenceTimeScaler execution completed successfully") + return result + + except Exception as e: + logger.error(f"Error in InferenceTimeScaler: {str(e)}") + logger.debug(traceback.format_exc()) + # Return a blank image and error message + error_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + error_json = json.dumps({"error": str(e)}) + return (error_tensor, error_tensor, error_json) + + def _execute_zero_order(self, model, vae, seed, steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, denoise, text_prompt_to_compare, + verifier_names, search_rounds, num_neighbors, lambda_threshold, view_top_k): + """ + Execute zero-order optimization search algorithm + """ + logger.info(f"Starting zero-order search with {search_rounds} rounds") + batch_inds = latent_image.get("batch_index", None) + pivot_noise = comfy.sample.prepare_noise(latent_image["samples"], seed, batch_inds) + pivot_included_in_pool = False + + # Initialize a pool to accumulate (candidate_tensor, raw_scores) pairs + candidate_pool = [] + + for round_idx in range(search_rounds): + # Generate neighbors around the pivot noise + neighbors = generate_neighbors(pivot_noise, threshold=lambda_threshold, num_neighbors=num_neighbors) + + # Prepend the pivot as the "zeroth" neighbor + pivot_expanded = pivot_noise.unsqueeze(1) # shape: [B,1,C,H,W] + neighbors = torch.cat([pivot_expanded, neighbors], dim=1) # shape: [B, num_neighbors+1, C,H,W] + + candidate_noises = [] + candidate_decoded_tensors = [] + + for i in range(num_neighbors + 1): + candidate_noises.append(neighbors[:, i]) # each candidate is of shape [B, C, H, W] + + # Decode all candidate noises for this round + for noise_candidate in candidate_noises: + samples, = common_ksampler( + model, seed, steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, denoise, + override_noise=noise_candidate + ) + decoded = vae.decode(samples["samples"]) + if decoded.ndim == 5: + b, c, h, w, d = decoded.shape + decoded = decoded.view(-1, h, w, d) + candidate_decoded_tensors.append(decoded[0:1]) + + # score the candidates + raw_scores = score_candidates(candidate_decoded_tensors, text_prompt_to_compare, verifier_names) + ranked_candidates = rank_from_raw_scores(raw_scores) + best_idx = ranked_candidates[0]["index"] + + for i, (tensor_img, candidate_raw_scores) in enumerate(zip(candidate_decoded_tensors, raw_scores)): + if i == 0: # pivot candidate + if not pivot_included_in_pool: + candidate_pool.append((tensor_img, candidate_raw_scores)) + pivot_included_in_pool = True + else: + candidate_pool.append((tensor_img, candidate_raw_scores)) + + # Rejecting a round if the pivot is best + if best_idx == 0: + # The original pivot outperformed all neighbors so we reject this round + pivot_noise = comfy.sample.prepare_noise(latent_image["samples"], seed + round_idx + 1000, batch_inds) + pivot_included_in_pool = False + else: + # Update the pivot to be the best candidate's noise from this round + pivot_noise = candidate_noises[best_idx] + pivot_included_in_pool = True + + # Final ranking from the candidate pool + pool_raw_scores = [scores for (_, scores) in candidate_pool] + ranked_pool = rank_from_raw_scores(pool_raw_scores) + # Reorder candidate_pool by the ranking results + candidate_pool_sorted = [candidate_pool[entry["index"]][0] for entry in ranked_pool] + top_k_tensors = candidate_pool_sorted[:view_top_k] + best_tensor = top_k_tensors[0] + + # Build ranking result list for JSON output + top_k_results = ranked_pool[:view_top_k] + + # Create grid of top-k images for visualization purposes + if len(top_k_tensors) == 0: + grid_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + else: + images_for_grid = [t.squeeze(0).permute(2, 0, 1) for t in top_k_tensors] + grid = make_grid(torch.stack(images_for_grid), nrow=len(images_for_grid)) + grid_image = grid.permute(1, 2, 0).unsqueeze(0) + + results_json = { + "results": top_k_results + } + scores_json_str = json.dumps(results_json, indent=2) + + return (best_tensor, grid_image, scores_json_str) + + def _execute_random_search(self, model, vae, seed, steps, cfg, sampler_name, scheduler, + positive, negative, latent_image, denoise, text_prompt_to_compare, + verifier_names, search_rounds, view_top_k): + """ + Execute random search algorithm + """ + logger.info(f"Starting random search with budget {search_rounds}") + candidate_tensors = [] + for i in range(search_rounds): new_seed = seed + i + 1 - samples, = common_ksampler(model, new_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise) - - candidate_images = vae.decode(samples["samples"]) - if len(candidate_images.shape) == 5: - candidate_images = candidate_images.reshape(-1, candidate_images.shape[-3], candidate_images.shape[-2], candidate_images.shape[-1]) - - candidate_image = candidate_images[0:1] # Using the first image of the batch. + samples, = common_ksampler(model, new_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise) + decoded = vae.decode(samples["samples"]) + if decoded.ndim == 5: + b, c, h, w, d = decoded.shape + decoded = decoded.view(-1, h, w, d) + single_candidate = decoded[0:1] + candidate_tensors.append(single_candidate) - # Convert from tensor to PIL Image - # Remove the batch dimension: [1, H, W, 3] -> [H, W, 3] - image_tensor = candidate_image.squeeze(0) - # Permute to [C, H, W] - image_tensor = image_tensor.permute(2, 0, 1) - # Convert to PIL Image - candidate_pil_image = ToPILImage()(image_tensor) + # scoring and ranking + raw_scores = score_candidates(candidate_tensors, text_prompt_to_compare, verifier_names) + final_ranked_results = rank_from_raw_scores(raw_scores) + top_k_results = final_ranked_results[:min(view_top_k, len(final_ranked_results))] + top_k_indices = [r["index"] for r in top_k_results] + top_k_tensors = [candidate_tensors[i] for i in top_k_indices] + best_tensor = candidate_tensors[final_ranked_results[0]["index"]] - # Score the candidate - combined_score, individual_scores = compute_combined_verifier_score(candidate_pil_image, text_prompt_to_compare, verifiers, weights) - - candidate_results.append((candidate_image, combined_score)) - - candidate_results.sort(key=lambda x: x[1], reverse=True) - top_candidates = candidate_results[:min(top_k, len(candidate_results))] - - images_list = [img for (img, score) in top_candidates] - scores_list = [score for (img, score) in top_candidates] - - # Prepare the images for grid display. - images_for_grid = [img.squeeze(0).permute(2, 0, 1) for img in images_list] - grid_image = make_grid(torch.stack(images_for_grid), nrow=top_k) - grid_image = grid_image.permute(1, 2, 0).unsqueeze(0) - - scores_json = json.dumps({"results": [{"score": score} for score in scores_list]}) - best_candidate = images_list[0] - return (best_candidate, grid_image, scores_json) + # Create grid of top-k images for visualization purposes + if len(top_k_tensors) == 0: + grid_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + else: + images_for_grid = [t.squeeze(0).permute(2, 0, 1) for t in top_k_tensors] + grid = make_grid(torch.stack(images_for_grid), nrow=len(images_for_grid)) + grid_image = grid.permute(1, 2, 0).unsqueeze(0) + + results_json = { + "results": top_k_results + } + scores_json_str = json.dumps(results_json, indent=2) + + return (best_tensor, grid_image, scores_json_str) +class LoadQwenVLMVerifier: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "qwen_verifier_id": (["Qwen/Qwen2.5-VL-7B-Instruct", None], { + "default": "Qwen/Qwen2.5-VL-7B-Instruct", + "tooltip": "Identifier for the Qwen VLM model." + }), + "device": ("STRING", { + "default": "cuda" if torch.cuda.is_available() else "cpu", + "tooltip": "Device to load the model onto." + }) + } + } + + RETURN_TYPES = ("QWN_VERIFIER",) + RETURN_NAMES = ("qwen_verifier_instance",) + FUNCTION = "execute" + CATEGORY = "InferenceTimeScaling" + DESCRIPTION = "Loads the Qwen VLM verifier model. Downloads it if necessary." + + def execute(self, qwen_verifier_id, device): + # Construct a local comfyui checkpoint path for the model + model_checkpoint = os.path.join(folder_paths.models_dir, "LLM", os.path.basename(qwen_verifier_id)) + if not os.path.exists(model_checkpoint): + snapshot_download( + repo_id=qwen_verifier_id, + local_dir=model_checkpoint, + local_dir_use_symlinks=False, + ) + verifier_instance = QwenVLMVerifier(qwen_verifier_id, device) + return (verifier_instance,) +class CLIPScoreVerifier: + def __init__(self, model, processor, device): + self.model = model + self.processor = processor + self.device = device + + def score(self, text_prompt, pil_img): + # Prepare inputs and move to the correct device + inputs = self.processor(text=[text_prompt], images=pil_img, return_tensors="pt", padding=True) + inputs = {k: v.to(self.device) for k, v in inputs.items()} + # Compute features + text_emb = self.model.get_text_features(inputs["input_ids"], attention_mask=inputs["attention_mask"]) + image_emb = self.model.get_image_features(pixel_values=inputs["pixel_values"]) + # Normalize and compute cosine similarity + text_emb = F.normalize(text_emb, dim=-1) + image_emb = F.normalize(image_emb, dim=-1) + return F.cosine_similarity(text_emb, image_emb).item() + + +class LoadCLIPScoreVerifier: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "clip_verifier_id": (["openai/clip-vit-base-patch32", "openai/clip-vit-large-patch14", None], { + "default": "openai/clip-vit-base-patch32", + "tooltip": "Identifier for the CLIP model." + }), + "device": ("STRING", { + "default": "cuda" if torch.cuda.is_available() else "cpu", + "tooltip": "Device to load the model onto." + }) + } + } + + RETURN_TYPES = ("CS_VERIFIER",) + RETURN_NAMES = ("clip_verifier_instance",) + FUNCTION = "execute" + CATEGORY = "InferenceTimeScaling" + DESCRIPTION = "Loads the CLIP verifier model and processor." + + def execute(self, clip_verifier_id, device): + clip_model = CLIPModel.from_pretrained(clip_verifier_id) + clip_processor = CLIPProcessor.from_pretrained(clip_verifier_id) + clip_model.to(device) + verifier_instance = CLIPScoreVerifier(clip_model, clip_processor, device) + return (verifier_instance,) + + +class ImageRewardVerifier: + def __init__(self, model, device): + self.model = model + self.device = device + + def score(self, text_prompt, pil_img): + return self.model.score(text_prompt, pil_img) + + +class LoadImageRewardVerifier: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "ir_verifier_id": (["ImageReward-v1.0", None], { + "default": "ImageReward-v1.0", + "tooltip": "Identifier for the ImageReward model." + }), + "device": ("STRING", { + "default": "cuda" if torch.cuda.is_available() else "cpu", + "tooltip": "Device to load the model onto." + }) + } + } + + RETURN_TYPES = ("IR_VERIFIER",) + RETURN_NAMES = ("image_reward_verifier_instance",) + FUNCTION = "execute" + CATEGORY = "InferenceTimeScaling" + DESCRIPTION = "Loads the ImageReward verifier model." + + def execute(self, ir_verifier_id, device): + ir_model = reward.load(ir_verifier_id) + try: + ir_model.to(device) + except Exception: + # Some ImageReward models may not support .to(device) directly + pass + verifier_instance = ImageRewardVerifier(ir_model, device) + return (verifier_instance,) # Node registration NODE_CLASS_MAPPINGS = { - "MaximKSampler": MaximKSampler + "InferenceTimeScaler": InferenceTimeScaler, + "LoadQwenVLMVerifier": LoadQwenVLMVerifier, + "LoadCLIPScoreVerifier": LoadCLIPScoreVerifier, + "LoadImageRewardVerifier": LoadImageRewardVerifier } - NODE_DISPLAY_NAME_MAPPINGS = { - "MaximKSampler": "MaximKSampler" + "InferenceTimeScaler": "Inference Time Scaler", + "LoadQwenVLMVerifier": "Load Qwen VLM Verifier", + "LoadCLIPScoreVerifier": "Load CLIPScore Verifier", + "LoadImageRewardVerifier": "Load ImageReward Verifier" } - -