LoadFlorenceModel node support new version PromptGen-v1.5 model

chflame163 · Sep 30, 2024 · 28c1a4f · 28c1a4f
1 parent 348539b
commit 28c1a4f
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 5 deletions.
diff --git a/py/florence2_ultra.py b/py/florence2_ultra.py
@@ -52,13 +52,36 @@ def load_model(version):
             model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
             processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
     except Exception as e:
-        log(f"Error loading model {version}: {str(e)}")
-        log("Attempting to load tokenizer instead of processor...")
         try:
             model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
             processor = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
         except Exception as e:
-            log(f"Error loading model or tokenizer: {str(e)}")
+            sys.path.append(model_path)
+            # Import the Florence modules
+            if version == 'large-PromptGen-v1.5':
+                from florence2_large.modeling_florence2 import Florence2ForConditionalGeneration
+                from florence2_large.configuration_florence2 import Florence2Config
+            elif version == 'base-PromptGen-v1.5':
+                from florence2_base_ft.modeling_florence2 import Florence2ForConditionalGeneration
+                from florence2_base_ft.configuration_florence2 import Florence2Config
+            else:
+                log(f"Error loading model or tokenizer: {str(e)}", message_type='error')
+                return (None, None)
+
+            attention = 'sdpa'
+            # Load the model configuration
+            model_config = Florence2Config.from_pretrained(model_path)
+            # Load the model
+            with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+                model = Florence2ForConditionalGeneration.from_pretrained(
+                    model_path,
+                    config=model_config,
+                    attn_implementation=attention,
+                    device_map=device
+                ).to(device)
+
+            processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+
 
     return (model.to(device), processor)
 

diff --git a/py/imagefunc.py b/py/imagefunc.py
@@ -30,7 +30,7 @@
 from skimage import img_as_float, img_as_ubyte
 import torchvision.transforms.functional as TF
 import torch.nn.functional as F
-from transformers import AutoModel, AutoProcessor, StoppingCriteria, StoppingCriteriaList, AutoModelForCausalLM
+from transformers import AutoModel, AutoProcessor, StoppingCriteria, StoppingCriteriaList, AutoModelForCausalLM, AutoTokenizer
 from colorsys import rgb_to_hsv
 import folder_paths
 import comfy.model_management

diff --git 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.70"
+version = "1.0.71"
 license = "MIT"
 dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "scikit_learn", "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"]