From a5e972af0522e8739b7c36df2a4ab0dd5af6f073 Mon Sep 17 00:00:00 2001
From: xingren23 <yexingren23@126.com>
Date: Sun, 10 Mar 2024 23:26:37 +0800
Subject: [PATCH] lazy load, optimize load time 0.5s to 0.0s

---
 py/imagefunc.py         | 3 ++-
 py/person_mask_Ultra.py | 6 +++---
 py/prompt_tagger.py     | 3 +--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/py/imagefunc.py b/py/imagefunc.py
index 46b543c..c15ae66 100644
--- a/py/imagefunc.py
+++ b/py/imagefunc.py
@@ -21,7 +21,6 @@ from functools import lru_cache
 from typing import Union, List
 from PIL import Image, ImageFilter, ImageChops, ImageDraw, ImageOps, ImageEnhance, ImageFont
 from skimage import img_as_float, img_as_ubyte
-from pymatting import fix_trimap, estimate_alpha_cf, estimate_foreground_ml
 from transformers import VitMatteImageProcessor, VitMatteForImageMatting
 import torchvision.transforms.functional as TF
 import torch.nn.functional as F
@@ -933,6 +932,7 @@ def image_beauty(image:Image, level:int=50) -> Image:
 
 
 def pixel_spread(image:Image, mask:Image) -> Image:
+    from pymatting import estimate_foreground_ml
     i1 = pil2tensor(image)
     if mask.mode != 'RGB':
         mask = mask.convert('RGB')
@@ -1095,6 +1095,7 @@ def get_a_person_mask_generator_model_path() -> str:
     return model_file_path
 
 def mask_edge_detail(image:torch.Tensor, mask:torch.Tensor, detail_range:int=8, black_point:float=0.01, white_point:float=0.99) -> torch.Tensor:
+    from pymatting import fix_trimap, estimate_alpha_cf
     d = detail_range * 5 + 1
     mask = pil2tensor(tensor2pil(mask).convert('RGB'))
     if not bool(d % 2):
diff --git a/py/person_mask_Ultra.py b/py/person_mask_Ultra.py
index 0c37074..dbc71d2 100644
--- a/py/person_mask_Ultra.py
+++ b/py/person_mask_Ultra.py
@@ -1,7 +1,6 @@
 from .imagefunc import *
 from functools import reduce
 import wget
-import mediapipe as mp
 import folder_paths
 from .segment_anything_func import *
 
@@ -43,7 +42,8 @@ class PersonMaskUltra:
     CATEGORY = '😺dzNodes/LayerMask'
     OUTPUT_NODE = True
 
-    def get_mediapipe_image(self, image: Image) -> mp.Image:
+    def get_mediapipe_image(self, image: Image):
+        import mediapipe as mp
         # Convert image to NumPy array
         numpy_image = np.asarray(image)
         image_format = mp.ImageFormat.SRGB
@@ -58,7 +58,7 @@ class PersonMaskUltra:
     def person_mask_ultra(self, images, face, hair, body, clothes,
                           accessories, background, confidence,
                           detail_range, black_point, white_point, process_detail):
-
+        import mediapipe as mp
         a_person_mask_generator_model_path = get_a_person_mask_generator_model_path()
         a_person_mask_generator_model_buffer = None
         with open(a_person_mask_generator_model_path, "rb") as f:
diff --git a/py/prompt_tagger.py b/py/prompt_tagger.py
index 9087481..e9d4f14 100644
--- a/py/prompt_tagger.py
+++ b/py/prompt_tagger.py
@@ -1,5 +1,4 @@
 from .imagefunc import *
-import google.generativeai as genai
 
 NODE_NAME = 'PromptTagger'
 
@@ -31,7 +30,7 @@ class PromptTagger:
     OUTPUT_NODE = True
 
     def prompt_tagger(self, image, api, token_limit, exclude_word, replace_with_word):
-
+        import google.generativeai as genai
         replace_with_word = replace_with_word.strip()
         exclude_word = exclude_word.strip()
         _image = tensor2pil(image).convert('RGB')