diff --git a/nodes.py b/nodes.py index 53c2f86..61fb1d1 100644 --- a/nodes.py +++ b/nodes.py @@ -102,11 +102,13 @@ class Florence2Run: [ 'region_caption', 'dense_region_caption', + 'region_proposal', 'caption', 'detailed_caption', 'more_detailed_caption', 'caption_to_phrase_grounding', - 'referring_expression_segmentation' + 'referring_expression_segmentation', + 'ocr' ], ), @@ -116,6 +118,7 @@ class Florence2Run: "keep_model_loaded": ("BOOLEAN", {"default": False}), "max_new_tokens": ("INT", {"default": 1024, "min": 1, "max": 4096}), "num_beams": ("INT", {"default": 3, "min": 1, "max": 64}), + "do_sample": ("BOOLEAN", {"default": True}), } } @@ -124,7 +127,8 @@ class Florence2Run: FUNCTION = "encode" CATEGORY = "Florence2" - def encode(self, image, text_input, florence2_model, task, fill_mask, keep_model_loaded=False, num_beams=3, max_new_tokens=1024): + def encode(self, image, text_input, florence2_model, task, fill_mask, keep_model_loaded=False, + num_beams=3, max_new_tokens=1024, do_sample=True): device = mm.get_torch_device() offload_device = mm.unet_offload_device() annotated_image_tensor = None @@ -136,26 +140,24 @@ class Florence2Run: colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red', 'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue'] - if task == 'region_caption': - prompt = "" - elif task == 'dense_region_caption': - prompt = '' - elif task == 'caption': - prompt = '' - elif task == 'detailed_caption': - prompt = '' - elif task == 'more_detailed_caption': - prompt = '' - elif task == 'caption_to_phrase_grounding': - prompt = '' - elif task == 'referring_expression_segmentation': - prompt = '' + prompts = { + 'region_caption': '', + 'dense_region_caption': '', + 'region_proposal': '', + 'caption': '', + 'detailed_caption': '', + 'more_detailed_caption': '', + 'caption_to_phrase_grounding': '', + 'referring_expression_segmentation': '', + 'ocr': '' + } + task_prompt = prompts.get(task, '') if (task!= 'referring_expression_segmentation' and task!= 'caption_to_phrase_grounding') and text_input: - raise ValueError("text_input is only supported for 'referring_expression_segmentation' and 'caption_to_phrase_grounding'") + raise ValueError("Text input (prompt) is only supported for 'referring_expression_segmentation' and 'caption_to_phrase_grounding'") if text_input is not None: - prompt = prompt + text_input + prompt = task_prompt + text_input image = image.permute(0, 3, 1, 2) @@ -171,12 +173,11 @@ class Florence2Run: input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=max_new_tokens, - do_sample=False, + do_sample=do_sample, num_beams=num_beams, ) results = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] - # cleanup the special tokens from the final list clean_results = str(results) clean_results = clean_results.replace('', '') @@ -188,14 +189,14 @@ class Florence2Run: else: out_results.append(clean_results) - if task == 'region_caption' or task == 'dense_region_caption' or task == 'caption_to_phrase_grounding': - parsed_answer = processor.post_process_generation(results, task="", image_size=(image_pil.width, image_pil.height)) + if task == 'region_caption' or task == 'dense_region_caption' or task == 'caption_to_phrase_grounding' or task == 'region_proposal': + parsed_answer = processor.post_process_generation(results, task=task_prompt, image_size=(image_pil.width, image_pil.height)) fig, ax = plt.subplots(figsize=(image_pil.width / 100, image_pil.height / 100), dpi=100) fig.subplots_adjust(left=0, right=1, top=1, bottom=0) ax.imshow(image_pil) - bboxes = parsed_answer['']['bboxes'] - labels = parsed_answer['']['labels'] + bboxes = parsed_answer[task_prompt]['bboxes'] + labels = parsed_answer[task_prompt]['labels'] # Loop through the bounding boxes and labels and add them to the plot for bbox, label in zip(bboxes, labels): # Create a Rectangle patch @@ -258,7 +259,7 @@ class Florence2Run: plt.close(fig) elif task == 'referring_expression_segmentation': - parsed_answer = processor.post_process_generation(results, task="", image_size=(image_pil.width, image_pil.height)) + parsed_answer = processor.post_process_generation(results, task=task_prompt, image_size=(image_pil.width, image_pil.height)) width, height = image_pil.size # Create a new black image mask_image = Image.new('RGB', (width, height), 'black')