From f2f8c16d86fa7eb4f25d6580381a06a7efc76daa Mon Sep 17 00:00:00 2001
From: chflame163 <chflame@163.com>
Date: Thu, 12 Sep 2024 00:01:02 +0800
Subject: [PATCH] Florence2Image2Prompt node support PromptGen and other fl2
 models task

---
 py/florence2_ultra.py | 21 +++++++++++++++++++--
 pyproject.toml        |  2 +-
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/py/florence2_ultra.py b/py/florence2_ultra.py
index 4a8b339..c9acb9c 100644
--- a/py/florence2_ultra.py
+++ b/py/florence2_ultra.py
@@ -26,7 +26,6 @@ fl2_model_repos = {
     "large-PromptGen-v1.5":"MiaoshouAI/Florence-2-large-PromptGen-v1.5"
 }
 
-
 def fixed_get_imports(filename) -> list[str]:
     """Workaround for FlashAttention"""
     if os.path.basename(filename) != "modeling_florence2.py":
@@ -283,6 +282,21 @@ def process_image(model, processor, image, task_prompt, max_new_tokens, num_beam
         output_results = {'bboxes': results[task_prompt].get('quad_boxes', []),
                           'labels': results[task_prompt].get('labels', [])}
         return output_results, output_image
+    # gokaygokay/Florence-2-SD3-Captioner task
+    elif task_prompt == 'description':
+        task_prompt = '<DESCRIPTION>'
+        result = run_example(model, processor, task_prompt, image, max_new_tokens, num_beams, do_sample)
+        return result[task_prompt], None
+    # MiaoshouAI/Florence-2-large-PromptGen-v1.5 task
+    elif task_prompt == 'generate tags(PromptGen 1.5)':
+        task_prompt = '<GENERATE_TAGS>'
+        result = run_example(model, processor, task_prompt, image, max_new_tokens, num_beams, do_sample)
+        return result[task_prompt], None
+    elif task_prompt == 'mixed caption(PromptGen 1.5)':
+        task_prompt = '<MIXED_CAPTION>'
+        result = run_example(model, processor, task_prompt, image, max_new_tokens, num_beams, do_sample)
+        return result[task_prompt], None
+
     else:
         return "", None  # Return empty string and None for unknown task prompts
 
@@ -490,6 +504,9 @@ class Florence2Image2Prompt:
             "caption",
             "detailed caption",
             "more detailed caption",
+            'description',
+            'generate tags(PromptGen 1.5)',
+            'mixed caption(PromptGen 1.5)',
             "object detection",
             "dense region caption",
             "region proposal",
@@ -499,7 +516,7 @@ class Florence2Image2Prompt:
             "region to category",
             "region to description",
             "OCR",
-            "OCR with region"
+            "OCR with region",
             ]
         return {
             "required": {
diff --git a/pyproject.toml b/pyproject.toml
index bc84b63..321eaad 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_layerstyle"
 description = "A set of nodes for ComfyUI it generate image like Adobe Photoshop's Layer Style. the Drop Shadow is first completed node, and follow-up work is in progress."
-version = "1.0.53"
+version = "1.0.54"
 license = "MIT"
 dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "opencv-contrib-python", "pymatting", "segment_anything", "timm", "addict", "yapf", "colour-science", "wget", "mediapipe", "loguru", "typer_config", "fastapi", "rich", "google-generativeai", "diffusers", "omegaconf", "tqdm", "transformers", "kornia", "image-reward", "ultralytics", "blend_modes", "blind-watermark", "qrcode", "pyzbar", "transparent-background", "huggingface_hub", "accelerate", "bitsandbytes", "torchscale", "wandb", "hydra-core", "psd-tools", "inference-cli[yolo-world]", "inference-gpu[yolo-world]", "onnxruntime"]