nodes/YRIKKA-ComfyUI-InferenceTimeScaling
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add configurable scoring criteria for Qwen VLM verifier

Browse Source 
- Add score_type parameter to QwenVLMVerifier
- Update documentation with new scoring options
- Add support for Qwen 3B and 72B models
- Update README with clearer verifier usage notes

Fixes #2
This commit is contained in:
MaximClouser
2025-02-27 19:48:01 +00:00
parent 4405870784
commit dbe9927e40
5 changed files with 69 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+22 -15
View File
@@ -23,6 +23,10 @@ A ComfyUI extension implementing "Inference-time scaling for diffusion models be
[View example workflows](workflows/)
## How It Works
It uses random search and zero-order search to explore different noises and a verifier ensemble (ImageReward, CLIP, and a VLM grader using Qwen 2.5 7b/3b/72b) to generate the best image. By exploring the noise space, it finds a better starting seed and produces images of higher quality and better prompt alignment than simply increasing denoising steps, with the tradeoff being increased time and compute during inference.
## Installation
### Prerequisites
@@ -79,15 +83,19 @@ This is the main node implementing the random search and zero-order optimization
- `vae`: (VAE) VAE model for decoding latents
- `view_top_k`: (INT) Number of top images to show in grid
- `search_algorithm`: Choice between "random" and "zero-order"
- `num_neighbors`: (INT) Number of neighbors per iteration in zero-order search (only used if search_algorithm is "zero-order")
- `lambda_threshold`: (FLOAT) Perturbation step size for zero-order search (only used if search_algorithm is "zero-order")
> [!IMPORTANT]
> The following parameters are **only used for zero-order search** and have no effect when using random search:
> - `num_neighbors`: (INT) Number of neighbors per iteration in zero-order search
> - `lambda_threshold`: (FLOAT) Perturbation step size for zero-order search
#### Optional Inputs:
- `loaded_clip_score_verifier`: (CS_VERIFIER) CLIP model for scoring
- `loaded_image_reward_verifier`: (IR_VERIFIER) ImageReward model
- `loaded_qwen_verifier`: (QWN_VERIFIER) Qwen VLM model
**Note:** At least one verifier must be included!
> [!NOTE]
> The verifiers are optional - you can choose which ones to use by connecting them to the node. However, at least one verifier must be connected for the node to function!
#### Outputs:
- `Best Image`: The highest-scoring generated image
@@ -103,6 +111,13 @@ Loads the Qwen VLM verifier model for image evaluation.
#### Inputs:
- `qwen_verifier_id`: Model identifier (default: "Qwen/Qwen2.5-VL-7B-Instruct")
- `device`: Device to load model on ("cuda" or "cpu")
- `score_type`: Type of score to return from the evaluation (default: "overall_score"). Options:
- `overall_score`: Weighted average of all aspects
- `accuracy_to_prompt`: How well the image matches the text description
- `creativity_and_originality`: Uniqueness and creative interpretation
- `visual_quality_and_realism`: Overall visual quality, detail, and realism
- `consistency_and_cohesion`: Internal consistency and natural composition
- `emotional_or_thematic_resonance`: How well the image captures the intended mood/theme
#### Outputs:
- `qwen_verifier_instance`: Loaded Qwen verifier instance
@@ -154,21 +169,13 @@ The model will be downloaded automatically on first use (you do not need to have
## Future Work
- [x] Enable configurable scoring criteria for Qwen VLM verifier
- Allow users to select specific aspects like visual quality, creativity, etc.
- Support individual aspect scoring
- [ ] Add batch processing support for image generation (performance optimization)
- [ ] Implement batched verification for multiple image-text pairs (speed optimization)
- [ ] Enable configurable scoring criteria for Qwen VLM verifier (currently only uses overall score)
- Allow users to select specific aspects like visual quality, creativity, etc.
- Support weighted combinations of multiple scoring criteria
- [ ] Add support for image-to-image and image+text conditioning to image models (currently only supports text-to-image models)
## Development
To install development dependencies:
```bash
cd inferencescale
pip install -e .[dev]
pre-commit install
```
## License
assets/qwen_verifier_node.jpeg
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 22 KiB

After

Width:  |  Height:  |  Size: 28 KiB

pyproject.toml
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "ComfyUI-InferenceTimeScaling"
version = "0.0.2"
version = "0.0.3"
description = "Inference-time techniques to enhance diffusion-based image generation quality through random search and zero-order optimization algorithms"
authors = [
{name = "Max Clouser", email = "max@yrikka.com"}
src/inferencescale/nodes.py
+15 -6
View File
@@ -149,7 +149,7 @@ def score_candidates(candidate_tensors: List[torch.Tensor], text_prompt: str, ve
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)
score = verifier.score(pil_img, text_prompt)
else:
logger.warning(f"Unknown verifier type: {verifier_name}")
continue
@@ -435,13 +435,19 @@ class LoadQwenVLMVerifier:
def INPUT_TYPES(cls):
return {
"required": {
"qwen_verifier_id": (["Qwen/Qwen2.5-VL-7B-Instruct", None], {
"qwen_verifier_id": (["Qwen/Qwen2.5-VL-7B-Instruct", "Qwen/Qwen2.5-VL-3B-Instruct", "Qwen/Qwen2.5-VL-72B-Instruct"], {
"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."
}),
"score_type": (["overall_score", "accuracy_to_prompt", "creativity_and_originality",
"visual_quality_and_realism", "consistency_and_cohesion",
"emotional_or_thematic_resonance"], {
"default": "overall_score",
"tooltip": "Type of score to return from the Qwen model evaluation."
})
}
}
@@ -452,7 +458,7 @@ class LoadQwenVLMVerifier:
CATEGORY = "InferenceTimeScaling"
DESCRIPTION = "Loads the Qwen VLM verifier model. Downloads it if necessary."
def execute(self, qwen_verifier_id, device):
def execute(self, qwen_verifier_id, device, score_type):
# 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):
@@ -461,7 +467,7 @@ class LoadQwenVLMVerifier:
local_dir=model_checkpoint,
local_dir_use_symlinks=False,
)
verifier_instance = QwenVLMVerifier(qwen_verifier_id, device)
verifier_instance = QwenVLMVerifier(model_checkpoint, device, score_type=score_type)
return (verifier_instance,)
@@ -489,7 +495,10 @@ class LoadCLIPScoreVerifier:
def INPUT_TYPES(cls):
return {
"required": {
"clip_verifier_id": (["openai/clip-vit-base-patch32", "openai/clip-vit-large-patch14", None], {
"clip_verifier_id": (["openai/clip-vit-base-patch32",
"openai/clip-vit-large-patch14",
"openai/clip-vit-base-patch16",
"openai/clip-vit-large-patch14-336"], {
"default": "openai/clip-vit-base-patch32",
"tooltip": "Identifier for the CLIP model."
}),
@@ -528,7 +537,7 @@ class LoadImageRewardVerifier:
def INPUT_TYPES(cls):
return {
"required": {
"ir_verifier_id": (["ImageReward-v1.0", None], {
"ir_verifier_id": (["ImageReward-v1.0"], {
"default": "ImageReward-v1.0",
"tooltip": "Identifier for the ImageReward model."
}),
src/inferencescale/qwen_verifier.py
+31 -54
View File
@@ -82,11 +82,25 @@ class Grading(BaseModel):
class QwenVLMVerifier():
def __init__(self, model_name, device='cpu'):
def __init__(self, model_name, device='cpu', score_type='overall_score'):
logger.info(f"Initializing QwenVLMVerifier with model {model_name} on device {device}")
self.model_name = model_name
self.device = device
self.dtype = torch.float16 if "cuda" in self.device else torch.float16
self.dtype = torch.float16
# Validate score_type
valid_score_types = {
'accuracy_to_prompt',
'creativity_and_originality',
'visual_quality_and_realism',
'consistency_and_cohesion',
'emotional_or_thematic_resonance',
'overall_score'
}
if score_type not in valid_score_types:
raise ValueError(f"Invalid score_type. Must be one of: {valid_score_types}")
self.score_type = score_type
self.load_model()
@@ -96,28 +110,14 @@ class QwenVLMVerifier():
min_pixels = 256 * 28 * 28
max_pixels = 1280 * 28 * 28
logger.debug(f"Loading model from {self.model_name}")
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
self.model_name,
torch_dtype=self.dtype,
device_map="auto",
low_cpu_mem_usage=True
)
logger.info("Model loaded successfully")
logger.debug("Loading processor")
processor = AutoProcessor.from_pretrained(
self.model_name, min_pixels=min_pixels, max_pixels=max_pixels
)
logger.info("Processor loaded successfully")
logger.debug("Initializing transformers vision")
self.qwen_model = transformers_vision(
self.model_name,
model_class=model.__class__,
model_class=Qwen2_5_VLForConditionalGeneration,
device=self.device,
model_kwargs={"torch_dtype": self.dtype},
processor_class=processor.__class__,
processor_class=AutoProcessor,
processor_kwargs={"min_pixels": min_pixels, "max_pixels": max_pixels}
)
logger.info("Transformers vision initialized")
@@ -125,8 +125,6 @@ class QwenVLMVerifier():
self.structured_qwen_generator = outlines.generate.json(self.qwen_model, Grading)
logger.info("Structured generator setup complete")
del model
del processor
torch.cuda.empty_cache()
gc.collect()
logger.info("Memory cleanup completed")
@@ -173,21 +171,20 @@ class QwenVLMVerifier():
raise
def get_overall_score(self, image, prompt: str, max_tokens: int = None, seed: int = 42) -> float:
# TODO Extend to handle any score key (instead of just overall score) as input
logger.info("Getting overall score")
def score(self, image, prompt: str, max_tokens: int = None, seed: int = 42) -> float:
logger.info(f"Getting {self.score_type} score")
try:
outputs = self.query_model(image, prompt, max_tokens, seed)
overall_score = outputs["overall_score"]["score"]
score = outputs[self.score_type]["score"]
if overall_score:
logger.debug(f"Overall score calculated: {overall_score}")
return float(overall_score)
if score is not None:
logger.debug(f"{self.score_type} score calculated: {score}")
return float(score)
logger.warning("Overall score not found in model output")
logger.warning(f"{self.score_type} score not found in model output")
return 0.0
except Exception as e:
logger.error(f"Error getting overall score: {str(e)}", exc_info=True)
logger.error(f"Error getting {self.score_type} score: {str(e)}", exc_info=True)
return 0.0
@@ -196,35 +193,15 @@ if __name__ == "__main__":
model_name = "Qwen/Qwen2.5-VL-7B-Instruct"
model = QwenVLMVerifier(model_name=model_name, device=device)
# model.load_model()
model = QwenVLMVerifier(model_name=model_name, device=device, score_type='visual_quality_and_realism')
image_path = "596F6DF4-2856-436E-A981-649ABFB15F1B.jpeg"
image = Image.open(image_path).convert("RGB")
test_prompt = "A red bird and a fish."
response = model.query_model(image, test_prompt)
print("Model Response:", response)
aspect_keys = [
"accuracy_to_prompt",
"creativity_and_originality",
"visual_quality_and_realism",
"consistency_and_cohesion",
"emotional_or_thematic_resonance"
]
scores = []
for key in aspect_keys:
if key in response and "score" in response[key]:
scores.append(response[key]["score"])
if scores:
average_score = sum(scores) / len(scores)
print("Average Score:", average_score)
else:
print("No scores found to average.")
visual_quality_score = model.score(image, test_prompt)
print(f"Visual quality score: {visual_quality_score}")
# model.to_device('cpu')
# model.to_device('cuda')
@@ -232,5 +209,5 @@ if __name__ == "__main__":
response = model.query_model(image, test_prompt)
print("Model Response:", response)
overall_score = model.get_overall_score(image, test_prompt)
overall_score = model.score(image, test_prompt)
print(f"Overall score: {overall_score}")