added the GeminiPromptEnhance node.

Author: zhangp365

README.md

@@ -39,6 +39,18 @@ This node, ColorCorrectOfUtils, is an extension of the original [ColorCorrect](h
 ## ModifyTextGender
 This node adjusts the text to describe the gender based on the input. If the gender input is 'M', the text will be adjusted to describe as male; if the gender input is 'F', it will be adjusted to describe as female.
 
+## GeminiPromptEnhance
+This node is designed to enhance the text description of the image, using the latest Gemini 2.0 flash model. It can add quality descriptors, lighting descriptions, scene descriptions, and skin descriptions to the text. and according to the gender input, can modifiy the content about gender.
+
+use this node, must get your free API key from Google AI Studio:
+   - Visit [Google AI Studio](https://aistudio.google.com/prompts/new_chat)
+   - Log in with your Google account
+   - Click on "Get API key" or go to settings
+   - Create a new API key
+   - Copy the API key for use in the node's input or gemini_config.json
+
+this code is original from https://github.com/ShmuelRonen/ComfyUI-Gemini_Flash_2.0_Exp, added new features. thanks to @ShmuelRonen.
+
 ## GenderControlOutput
 This node determines the output based on the input gender. If the gender input is 'M', it will output male-specific text, float, and integer values. If the gender input is 'F', it will output female-specific text, float, and integer values.
 

py/node_gemini_enhance_prompte.py

@@ -0,0 +1,199 @@
+# this code is original from https://github.com/ShmuelRonen/ComfyUI-Gemini_Flash_2.0_Exp, added cache and gender support
+import os
+import json
+import google.generativeai as genai
+from contextlib import contextmanager
+from collections import OrderedDict
+import folder_paths
+import logging
+import yaml
+logger = logging.getLogger(__name__)
+
+config_dir = os.path.join(folder_paths.base_path, "config")
+if not os.path.exists(config_dir):
+    os.makedirs(config_dir)
+
+
+def get_config():
+    try:
+        config_path = os.path.join(config_dir, 'gemini_config.yml')
+        with open(config_path, 'r') as f:  
+            config = yaml.load(f, Loader=yaml.FullLoader)
+        return config
+    except:
+        return {}
+
+def save_config(config):
+    config_path = os.path.join(config_dir, 'gemini_config.yml')
+    with open(config_path, 'w') as f:
+        yaml.dump(config, f, indent=4)
+
+@contextmanager
+def temporary_env_var(key: str, new_value):
+    old_value = os.environ.get(key)
+    if new_value is not None:
+        os.environ[key] = new_value
+    elif key in os.environ:
+        del os.environ[key]
+    try:
+        yield
+    finally:
+        if old_value is not None:
+            os.environ[key] = old_value
+        elif key in os.environ:
+            del os.environ[key]
+
+class LRUCache(OrderedDict):
+    def __init__(self, capacity):
+        super().__init__()
+        self.capacity = capacity
+
+    def get(self, key):
+        if key not in self:
+            return None
+        self.move_to_end(key)
+        return self[key]
+
+    def put(self, key, value):
+        if key in self:
+            self.move_to_end(key)
+        self[key] = value
+        if len(self) > self.capacity:
+            self.popitem(last=False)
+
+class GeminiPromptEnhance:
+    def __init__(self, api_key=None, proxy=None):
+        config = get_config()
+        self.api_key = api_key or config.get("GEMINI_API_KEY")
+        self.proxy = proxy or config.get("PROXY")
+        self.cache_size = 500  # 缓存最大条数
+        self.cache_file = os.path.join(config_dir, 'prompt_cache_gemini.yml')
+        self.cache = LRUCache(self.cache_size)
+        self.load_cache()
+        if self.api_key is not None:
+            self.configure_genai()
+
+    def load_cache(self):
+        try:
+            if os.path.exists(self.cache_file):
+                with open(self.cache_file, 'r', encoding='utf-8') as f:
+                    cache_data = yaml.load(f, Loader=yaml.FullLoader)
+                    # 重新创建LRU缓存
+                    for k, v in cache_data.items():
+                        self.cache.put(k, v)
+        except Exception as e:
+            logger.error(f"加载缓存出错: {str(e)}")
+            self.cache = LRUCache(self.cache_size)
+
+    def save_cache(self):
+        try:
+            with open(self.cache_file, 'w', encoding='utf-8') as f:
+                yaml.dump(dict(self.cache), f, indent=4)
+        except Exception as e:
+            logger.error(f"保存缓存出错: {str(e)}")
+
+    def configure_genai(self):
+        genai.configure(api_key=self.api_key, transport='rest')
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        default_prompt = "Edit and enhance the text description of the image. \nAdd quality descriptors, like 'A high-quality photo, an 8K photo.' \nAdd lighting descriptions based on the scene, like 'The lighting is natural and bright, casting soft shadows.' \nAdd scene descriptions according to the context, like 'The overall mood is serene and peaceful.' \nIf a person is in the scene, include a description of the skin, such as 'natural skin tones and ensure the skin appears realistic with clear, fine details.' \n\nOnly output the result of the text, no others.\nthe text is:"
+
+        return {
+            "required": {
+                "prompt": ("STRING", {"default": default_prompt, "multiline": True}),
+            },
+            "optional": {
+                "text_input": ("STRING", {"default": "", "multiline": True}),
+                "api_key": ("STRING", {"default": ""}),
+                "proxy": ("STRING", {"default": ""}),
+                "max_output_tokens": ("INT", {"default": 8192, "min": 1, "max": 8192}),
+                "temperature": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.1}),
+                "gender": (["","M", "F"], {"default": ""}),
+                "enabled": ("BOOLEAN", {"default": True}),                
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("generated_content",)
+    FUNCTION = "generate_content"
+    CATEGORY = "utils/text"
+
+    def prepare_content(self, prompt, text_input, gender=""):
+        if gender == "M":
+            prompt = "edit and enhance the text content according to male gender. if there is a female, must change the text to describe as male.\n" + prompt
+        elif gender == "F":
+            prompt = "edit and enhance the text content according to female gender. if there is a male, must change the text to describe as female.\n" + prompt
+            
+        text_content = prompt if not text_input else f"{prompt}\n{text_input}"
+        return [{"text": text_content}]
+
+    def generate_content(self, prompt, text_input=None, api_key="", proxy="",
+                        max_output_tokens=8192, temperature=0.4, gender="", enabled=True):
+        if not enabled:
+            return (text_input,)
+        
+        # 生成缓存键
+        cache_key = f"{text_input or ''}_{gender}"
+        
+        # 检查缓存
+        cached_result = self.cache.get(cache_key)
+        if cached_result is not None:
+            return (cached_result,)
+
+        # Set all safety settings to block_none by default
+        safety_settings = [
+            {"category": "harassment", "threshold": "NONE"},
+            {"category": "hate_speech", "threshold": "NONE"},
+            {"category": "sexually_explicit", "threshold": "NONE"},
+            {"category": "dangerous_content", "threshold": "NONE"},
+            {"category": "civic", "threshold": "NONE"}
+        ]
+
+        # Only update API key if explicitly provided in the node
+        if api_key.strip():
+            self.api_key = api_key
+            save_config({"GEMINI_API_KEY": self.api_key, "PROXY": self.proxy})
+            self.configure_genai()
+        
+        # Only update proxy if explicitly provided in the node    
+        if proxy.strip():
+            self.proxy = proxy
+            save_config({"GEMINI_API_KEY": self.api_key, "PROXY": self.proxy})
+
+        if not self.api_key:
+            raise ValueError("API key not found in config.json or node input")
+
+        model_name = 'models/gemini-2.0-flash-exp'
+        model = genai.GenerativeModel(model_name)
+
+        # Apply fixed safety settings to the model
+        model.safety_settings = safety_settings
+
+        generation_config = genai.types.GenerationConfig(
+            max_output_tokens=max_output_tokens,
+            temperature=temperature
+        )
+
+        with temporary_env_var('HTTP_PROXY', self.proxy), temporary_env_var('HTTPS_PROXY', self.proxy):
+            try:           
+                content_parts = self.prepare_content(prompt, text_input, gender)
+                response = model.generate_content(content_parts, generation_config=generation_config)
+                generated_content = response.text
+                
+                # 更新缓存
+                self.cache.put(cache_key, generated_content)
+                self.save_cache()
+                
+            except Exception as e:
+                generated_content = f"Error: {str(e)}"
+        
+        return (generated_content,)
+        
+NODE_CLASS_MAPPINGS = {
+    "GeminiPromptEnhance": GeminiPromptEnhance,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "GeminiPromptEnhance": "Gemini prompt enhance",
+}

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-utils-nodes"
 description = "Nodes:LoadImageWithSwitch, ImageBatchOneOrMore, ModifyTextGender, GenderControlOutput, ImageCompositeMaskedWithSwitch, ImageCompositeMaskedOneByOne, ColorCorrectOfUtils, SplitMask, MaskFastGrow, CheckpointLoaderSimpleWithSwitch, ImageResizeTo8x, MatchImageRatioToPreset, UpscaleImageWithModelIfNeed, MaskFromFaceModel, MaskCoverFourCorners, DetectorForNSFW, DeepfaceAnalyzeFaceAttributes etc."
-version = "1.2.6"
+version = "1.2.7"
 license = { file = "LICENSE" }
 dependencies = []
 

requirements.txt

@@ -7,4 +7,7 @@ onnxruntime>=1.19.2
 # DeepfaceAnalyzeFaceAttributes
 deepface==0.0.93
 ultralytics
-tf-keras
+tf-keras
+
+# Gemini_prompt_enhance nod
+google-generativeai>0.4.1