ComfyUI-Gemini/GeminiAPINode.py

import os
import io
import json
import requests
import torch
import google.generativeai as genai
from io import BytesIO
from PIL import Image

p = os.path.dirname(os.path.realpath(__file__))

def get_gemini_api_key():
    try:
        config_path = os.path.join(p, 'config.json')
        with open(config_path, 'r') as f:  
            config = json.load(f)
        api_key = config["GEMINI_API_KEY"]
    except:
        print("出错啦 Error: API key is required")
        return ""
    return api_key

class Gemini_API_Zho:

    def __init__(self, api_key=None):
        self.api_key = api_key
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "What is the meaning of life?", "multiline": True}),
                "model_name": (["gemini-pro", "gemini-1.5-flash", "gemini-1.5-pro-latest"],),
                "stream": ("BOOLEAN", {"default": False}),
                "api_key": ("STRING", {"default": ""})  # Add api_key as an input
            },
            "optional": {
                "image": ("IMAGE",),  
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "generate_content"

    CATEGORY = "Zho模块组/✨Gemini"

    def tensor_to_image(self, tensor):
        # 确保张量是在CPU上
        tensor = tensor.cpu()
    
        # 将张量数据转换为0-255范围并转换为整数
        # 这里假设张量已经是H x W x C格式
        image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
    
        # 创建PIL图像
        image = Image.fromarray(image_np, mode='RGB')
        return image

    def generate_content(self, prompt, model_name, stream, api_key, image=None):
        if api_key:
            self.api_key = api_key
            genai.configure(api_key=self.api_key,transport='rest')
        if not self.api_key:
            raise ValueError("API key is required")

        model = genai.GenerativeModel(model_name)

        if model_name == 'gemini-pro':
            if stream:
                response = model.generate_content(prompt, stream=True)
                textoutput = "\n".join([chunk.text for chunk in response])
            else:
                response = model.generate_content(prompt)
                textoutput = response.text
        
        if model_name == 'gemini-1.5-flash':
            if image == None:
                raise ValueError("gemini-1.5-flash needs image")
            else:
                # 转换图像
                pil_image = self.tensor_to_image(image)

                # 直接使用PIL图像
                if stream:
                    response = model.generate_content([prompt, pil_image], stream=True)
                    textoutput = "\n".join([chunk.text for chunk in response])
                else:
                    response = model.generate_content([prompt, pil_image])
                    textoutput = response.text

        if model_name == 'gemini-1.5-pro-latest':
            if image == None:
                if stream:
                    response = model.generate_content(prompt, stream=True)
                    textoutput = "\n".join([chunk.text for chunk in response])
                else:
                    response = model.generate_content(prompt)
                    textoutput = response.text
            else:
                # 转换图像
                pil_image = self.tensor_to_image(image)

                # 直接使用PIL图像
                if stream:
                    response = model.generate_content([prompt, pil_image], stream=True)
                    textoutput = "\n".join([chunk.text for chunk in response])
                else:
                    response = model.generate_content([prompt, pil_image])
                    textoutput = response.text
        
        return (textoutput,)


class Gemini_API_Vsion_ImgURL_Zho:

    def __init__(self, api_key=None):
        self.api_key = api_key
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "Describe this image", "multiline": True}),
                "image_url": ("STRING", {"default": ""}),
                "model_name": (["gemini-1.5-flash", "gemini-1.5-pro-latest"],),
                "stream": ("BOOLEAN", {"default": False}),
                "api_key": ("STRING", {"default": ""})  # Add api_key as an input
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "generate_content"

    CATEGORY = "Zho模块组/✨Gemini"

    def generate_content(self, prompt, model_name, stream, api_key, image_url):
        if api_key:
            self.api_key = api_key
            genai.configure(api_key=self.api_key,transport='rest')
        if not self.api_key:
            raise ValueError("API key is required")

        # Load the image from the URL
        response = requests.get(image_url)
        if response.status_code != 200:
            raise ValueError("Failed to load image from URL")
        img = Image.open(BytesIO(response.content))

        model = genai.GenerativeModel(model_name)

        if stream:
            response = model.generate_content([prompt, img], stream=True)
            textoutput = "\n".join([chunk.text for chunk in response])
        else:
            response = model.generate_content([prompt, img])
            textoutput = response.text
        
        return (textoutput,)

#chat
class Gemini_API_Chat_Zho:

    def __init__(self, api_key=None):
        self.api_key = api_key
        self.chat = None  # 初始化时，聊天实例为空
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "What is the meaning of life?", "multiline": True}),
                "model_name": (["gemini-pro", "gemini-1.5-pro-latest"],),
                "api_key": ("STRING", {"default": ""})  # Add api_key as an input
            },
            "optional": {
                "image": ("IMAGE",),  
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("response",)
    FUNCTION = "generate_chat"

    CATEGORY = "Zho模块组/✨Gemini"
    
    def generate_chat(self, prompt, model_name, api_key):
        if api_key:
            self.api_key = api_key
            genai.configure(api_key=self.api_key,transport='rest')
        if not self.api_key:
            raise ValueError("API key is required")

        if not self.chat:
            model = genai.GenerativeModel(model_name)
            self.chat = model.start_chat(history=[])

        if model_name == 'gemini-pro':
            response = self.chat.send_message(prompt)
            textoutput = response.text
            chat_history = self.format_chat_history(self.chat)

        if model_name == 'gemini-1.5-pro-latest':
            if image == None:
                response = self.chat.send_message(prompt)
                textoutput = response.text
                chat_history = self.format_chat_history(self.chat)
            else:
                # 转换图像
                pil_image = self.tensor_to_image(image)

                response = self.chat.send_message([prompt, pil_image])
                textoutput = response.text
                chat_history = self.format_chat_history(self.chat)
        
        return (chat_history,)

    def format_chat_history(self, chat):
        formatted_history = []
        for message in chat.history:
            formatted_message = f"{message.role}: {message.parts[0].text}"
            formatted_history.append(formatted_message)
            formatted_history.append("-" * 40)  # 添加分隔线
        return "\n".join(formatted_history)


class Gemini_API_S_Zho:

    def __init__(self):
        self.api_key = get_gemini_api_key()
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "What is the meaning of life?", "multiline": True}),
                "model_name": (["gemini-pro", "gemini-1.5-flash", "gemini-1.5-pro-latest"],),
                "stream": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "image": ("IMAGE",),  
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "generate_content"

    CATEGORY = "Zho模块组/✨Gemini"

    def tensor_to_image(self, tensor):
        # 确保张量是在CPU上
        tensor = tensor.cpu()
    
        # 将张量数据转换为0-255范围并转换为整数
        # 这里假设张量已经是H x W x C格式
        image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
    
        # 创建PIL图像
        image = Image.fromarray(image_np, mode='RGB')
        return image

    def generate_content(self, prompt, model_name, stream, image=None):
        if not self.api_key:
            raise ValueError("API key is required")

        model = genai.GenerativeModel(model_name)

        if model_name == 'gemini-pro':
            if stream:
                response = model.generate_content(prompt, stream=True)
                textoutput = "\n".join([chunk.text for chunk in response])
            else:
                response = model.generate_content(prompt)
                textoutput = response.text
        
        if model_name == 'gemini-1.5-flash':
            if image == None:
                raise ValueError("gemini-1.5-flash needs image")
            else:
                # 转换图像
                pil_image = self.tensor_to_image(image)

                # 直接使用PIL图像
                if stream:
                    response = model.generate_content([prompt, pil_image], stream=True)
                    textoutput = "\n".join([chunk.text for chunk in response])
                else:
                    response = model.generate_content([prompt, pil_image])
                    textoutput = response.text

        if model_name == 'gemini-1.5-pro-latest':
            if image == None:
                if stream:
                    response = model.generate_content(prompt, stream=True)
                    textoutput = "\n".join([chunk.text for chunk in response])
                else:
                    response = model.generate_content(prompt)
                    textoutput = response.text
            else:
                # 转换图像
                pil_image = self.tensor_to_image(image)

                # 直接使用PIL图像
                if stream:
                    response = model.generate_content([prompt, pil_image], stream=True)
                    textoutput = "\n".join([chunk.text for chunk in response])
                else:
                    response = model.generate_content([prompt, pil_image])
                    textoutput = response.text
        
        return (textoutput,)


class Gemini_API_S_Vsion_ImgURL_Zho:

    def __init__(self):
        self.api_key = get_gemini_api_key()
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "Describe this image", "multiline": True}),
                "image_url": ("STRING", {"default": ""}),
                "model_name": (["gemini-1.5-flash", "gemini-1.5-pro-latest"],),
                "stream": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "generate_content"

    CATEGORY = "Zho模块组/✨Gemini"

    def generate_content(self, prompt, model_name, stream, image_url):
        if not self.api_key:
            raise ValueError("API key is required")

        # Load the image from the URL
        response = requests.get(image_url)
        if response.status_code != 200:
            raise ValueError("Failed to load image from URL")
        img = Image.open(BytesIO(response.content))

        model = genai.GenerativeModel(model_name)

        if stream:
            response = model.generate_content([prompt, img], stream=True)
            textoutput = "\n".join([chunk.text for chunk in response])
        else:
            response = model.generate_content([prompt, img])
            textoutput = response.text
        
        return (textoutput,)


#chat
class Gemini_API_S_Chat_Zho:

    def __init__(self):
        self.api_key = get_gemini_api_key()
        self.chat = None  # 初始化时，聊天实例为空
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "What is the meaning of life?", "multiline": True}),
                "model_name": (["gemini-pro", "gemini-1.5-pro-latest"],),
            },
            "optional": {
                "image": ("IMAGE",),  
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("response",)
    FUNCTION = "generate_chat"

    CATEGORY = "Zho模块组/✨Gemini"

    def tensor_to_image(self, tensor):
        # 确保张量是在CPU上
        tensor = tensor.cpu()
    
        # 将张量数据转换为0-255范围并转换为整数
        # 这里假设张量已经是H x W x C格式
        image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
    
        # 创建PIL图像
        image = Image.fromarray(image_np, mode='RGB')
        return image
    
    def generate_chat(self, prompt, model_name):
        if not self.api_key:
            raise ValueError("API key is required")

        if not self.chat:
            model = genai.GenerativeModel(model_name)
            self.chat = model.start_chat(history=[])

        if model_name == 'gemini-pro':
            response = self.chat.send_message(prompt)
            textoutput = response.text
            chat_history = self.format_chat_history(self.chat)

        if model_name == 'gemini-1.5-pro-latest':
            if image == None:
                response = self.chat.send_message(prompt)
                textoutput = response.text
                chat_history = self.format_chat_history(self.chat)
            else:
                # 转换图像
                pil_image = self.tensor_to_image(image)

                response = self.chat.send_message([prompt, pil_image])
                textoutput = response.text
                chat_history = self.format_chat_history(self.chat)
        
        return (chat_history,)

    def format_chat_history(self, chat):
        formatted_history = []
        for message in chat.history:
            formatted_message = f"{message.role}: {message.parts[0].text}"
            formatted_history.append(formatted_message)
            formatted_history.append("-" * 40)  # 添加分隔线
        return "\n".join(formatted_history)


# System instructions
class Gemini_15P_API_S_Advance_Zho:

    def __init__(self):
        self.api_key = get_gemini_api_key()
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "What is the meaning of life?", "multiline": True}),
                "system_instruction": ("STRING", {"default": "You are creating a prompt for Stable Diffusion to generate an image. First step: describe this image, then put description into text. Second step: generate a text prompt for %s based on first step.  Only respond with the prompt itself, but embellish it as needed but keep it under 80 tokens.", "multiline": True}),
                "model_name": (["gemini-1.5-pro-latest"],),
                "stream": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "image": ("IMAGE",),  
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "generate_content"

    CATEGORY = "Zho模块组/✨Gemini"

    def tensor_to_image(self, tensor):
        # 确保张量是在CPU上
        tensor = tensor.cpu()
    
        # 将张量数据转换为0-255范围并转换为整数
        # 这里假设张量已经是H x W x C格式
        image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
    
        # 创建PIL图像
        image = Image.fromarray(image_np, mode='RGB')
        return image

    def generate_content(self, prompt, system_instruction, model_name, stream, image=None):
        if not self.api_key:
            raise ValueError("API key is required")

        model = genai.GenerativeModel(model_name, system_instruction=system_instruction)

        if image == None:
            if stream:
                response = model.generate_content(prompt, stream=True)
                textoutput = "\n".join([chunk.text for chunk in response])
            else:
                response = model.generate_content(prompt)
                textoutput = response.text
        else:
            # 转换图像
            pil_image = self.tensor_to_image(image)

            # 直接使用PIL图像
            if stream:
                response = model.generate_content([prompt, pil_image], stream=True)
                textoutput = "\n".join([chunk.text for chunk in response])
            else:
                response = model.generate_content([prompt, pil_image])
                textoutput = response.text

        return (textoutput,)


class Gemini_15P_API_S_Chat_Advance_Zho:

    def __init__(self):
        self.api_key = get_gemini_api_key()
        self.chat = None  # 初始化时，聊天实例为空
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"default": "What is the meaning of life?", "multiline": True}),
                "system_instruction": ("STRING", {"default": "You are creating a prompt for Stable Diffusion to generate an image. First step: describe this image, then put description into text. Second step: generate a text prompt for %s based on first step.  Only respond with the prompt itself, but embellish it as needed but keep it under 80 tokens.", "multiline": True}),
                "model_name": (["gemini-1.5-pro-latest"],),
            },
            "optional": {
                "image": ("IMAGE",),  
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("response",)
    FUNCTION = "generate_chat"

    CATEGORY = "Zho模块组/✨Gemini"
    
    def tensor_to_image(self, tensor):
        # 确保张量是在CPU上
        tensor = tensor.cpu()
    
        # 将张量数据转换为0-255范围并转换为整数
        # 这里假设张量已经是H x W x C格式
        image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
    
        # 创建PIL图像
        image = Image.fromarray(image_np, mode='RGB')
        return image

    def generate_chat(self, prompt, system_instruction, model_name, image=None):
        if not self.api_key:
            raise ValueError("API key is required")

        if not self.chat:
            model = genai.GenerativeModel(model_name, system_instruction=system_instruction)
            self.chat = model.start_chat(history=[])

        if model_name == 'gemini-1.5-pro-latest':
            if image == None:
                response = self.chat.send_message(prompt)
                textoutput = response.text
                chat_history = self.format_chat_history(self.chat)
            else:
                # 转换图像
                pil_image = self.tensor_to_image(image)

                response = self.chat.send_message([prompt, pil_image])
                textoutput = response.text
                chat_history = self.format_chat_history(self.chat)
        
        return (chat_history,)

    def format_chat_history(self, chat):
        formatted_history = []
        for message in chat.history:
            formatted_message = f"{message.role}: {message.parts[0].text}"
            formatted_history.append(formatted_message)
            formatted_history.append("-" * 40)  # 添加分隔线
        return "\n".join(formatted_history)


# File API
class Gemini_FileUpload_API_S_Zho:

    def __init__(self):
        self.api_key = get_gemini_api_key()
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "file": ("STRING", {"default": "./sample.mp3", "multiline": False}),
            }
        }

    RETURN_TYPES = ("FILE",)
    RETURN_NAMES = ("file",)
    FUNCTION = "file_upload"

    CATEGORY = "Zho模块组/✨Gemini"

    def file_upload(self, file):
        if not self.api_key:
            raise ValueError("API key is required")

        your_file = genai.upload_file(file)

        return [your_file]


class Gemini_File_API_S_Zho:

    def __init__(self):
        self.api_key = get_gemini_api_key()
        if self.api_key is not None:
            genai.configure(api_key=self.api_key,transport='rest')

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "file": ("FILE",),
                "prompt": ("STRING", {"default": "Listen carefully to the following audio file. Provide a brief summary.", "multiline": True}),
                "model_name": (["gemini-1.5-pro-latest"],),
                "stream": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "generate_content"

    CATEGORY = "Zho模块组/✨Gemini"

    def generate_content(self, prompt, model_name, stream, file):
        if not self.api_key:
            raise ValueError("API key is required")

        model = genai.GenerativeModel(model_name)

        if stream:
            response = model.generate_content([prompt, file], stream=True)
            textoutput = "\n".join([chunk.text for chunk in response])
        else:
            response = model.generate_content([prompt, file])
            textoutput = response.text

        return (textoutput,)


class ConcatText_Zho:

    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text_1": ("STRING", {"multiline": True}),
                "text_2": ("STRING", {"multiline": True}),
                # 可以根据需要添加更多的文本输入
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)

    FUNCTION = "concat_texts"

    CATEGORY = "Zho模块组/✨Gemini"

    def concat_texts(self, **kwargs):
        # 将所有输入的文本合并为一个以逗号分隔的字符串
        texts = [kwargs[key] for key in kwargs if key.startswith('text')]
        combined_text = ', '.join(texts)
        return (combined_text,)


# DisplayText node is forked from AlekPet，thanks to AlekPet！
class DisplayText_Zho:
    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(s):

        return {
            "required": {        
                "text": ("STRING", {"forceInput": True}),     
                },
            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
            }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    OUTPUT_NODE = True
    FUNCTION = "display_text"

    CATEGORY = "Zho模块组/✨Gemini"

    def display_text(self, text, prompt=None, extra_pnginfo=None):
        return {"ui": {"string": [text,]}, "result": (text,)}


NODE_CLASS_MAPPINGS = {
    "Gemini_API_Zho": Gemini_API_Zho,
    "Gemini_API_Vsion_ImgURL_Zho": Gemini_API_Vsion_ImgURL_Zho,
    "Gemini_API_Chat_Zho": Gemini_API_Chat_Zho,
    "Gemini_API_S_Zho": Gemini_API_S_Zho,
    "Gemini_API_S_Vsion_ImgURL_Zho": Gemini_API_S_Vsion_ImgURL_Zho,
    "Gemini_API_S_Chat_Zho": Gemini_API_S_Chat_Zho,
    "Gemini_15P_API_S_Advance_Zho": Gemini_15P_API_S_Advance_Zho,
    "Gemini_15P_API_S_Chat_Advance_Zho": Gemini_15P_API_S_Chat_Advance_Zho,
    "Gemini_FileUpload_API_S_Zho": Gemini_FileUpload_API_S_Zho,
    "Gemini_File_API_S_Zho": Gemini_File_API_S_Zho,
    "ConcatText_Zho": ConcatText_Zho,
    "DisplayText_Zho": DisplayText_Zho
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Gemini_API_Zho": "✨Gemini_API_Zho",
    "Gemini_API_Vsion_ImgURL_Zho": "✨Gemini_API_Vsion_ImgURL_Zho",
    "Gemini_API_Chat_Zho": "✨Gemini_API_Chat_Zho",
    "Gemini_API_S_Zho": "㊙️Gemini_Zho",
    "Gemini_API_S_Vsion_ImgURL_Zho": "㊙️Gemini_ImgURL_Zho",
    "Gemini_API_S_Chat_Zho": "㊙️Gemini_Chat_Zho",
    "Gemini_15P_API_S_Advance_Zho": "🆕Gemini_15P_Advance_Zho",
    "Gemini_15P_API_S_Chat_Advance_Zho": "🆕Gemini_15P_Chat_Advance_Zho",
    "Gemini_FileUpload_API_S_Zho": "📄Gemini_FileUpload_Zho",
    "Gemini_File_API_S_Zho": "📄Gemini_File_Zho",
    "ConcatText_Zho": "✨ConcatText_Zho",
    "DisplayText_Zho": "✨DisplayText_Zho"
}