From 1797f9c810aa93dfd38253d2b27cbe6c347d7238 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Mon, 8 Apr 2024 09:44:09 +0800
Subject: [PATCH] Update ChatGLM4_API_Node.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

去掉直接使用URL的方法，改用base64编码的方式与服务器交互
---
 ChatGLM4_API_Node.py | 44 +++++++++++++++++++++++++++++++++++---------
 1 file changed, 35 insertions(+), 9 deletions(-)

diff --git a/ChatGLM4_API_Node.py b/ChatGLM4_API_Node.py
index dda4e2b..10bc467 100644
--- a/ChatGLM4_API_Node.py
+++ b/ChatGLM4_API_Node.py
@@ -4,6 +4,11 @@ import requests
 import time
 import jwt
 import re
+import tempfile
+import random
+from PIL import Image
+import base64
+import cv2
 
 p = os.path.dirname(os.path.realpath(__file__))  
 
@@ -100,7 +105,7 @@ class ZhipuaiApi_Img:
         return {
             "required": {
                 "prompt": ("STRING", {"default": "Describe this image", "multiline": True}),
-                "img_url": ("STRING", {"default": "https://", "multiline": True}),
+                "image": ("IMAGE",),
             }
         }
 
@@ -109,13 +114,19 @@ class ZhipuaiApi_Img:
     FUNCTION = "zhipuai_img2txt_api"
     CATEGORY = "ChatGlm_Api"
 
-    def zhipuai_img2txt_api(self, prompt, img_url):
+    def tensor_to_image(self, tensor):
+        tensor = tensor.cpu()
+        image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
+        image = Image.fromarray(image_np, mode='RGB')
+        return image
+
+    def zhipuai_img2txt_api(self, prompt, image):
         if not self.api_key:
             raise ValueError("API key is required")
         if prompt == None:
             raise ValueError("need prompt")
-        if img_url == None:
-            raise ValueError("Needs a Internet img url")
+        if image == None:
+            raise ValueError("Needs img")
         else:
             prompt_img = 'in English'
             prompt = ''.join([prompt, prompt_img])
@@ -127,16 +138,28 @@ class ZhipuaiApi_Img:
             client_key = str(get_zpai_api_key())
 
             token = generate_token(apikey_s=client_key, exp_second=exp_seconds)
-            url_img = f"{img_url}"
+
+            pil_image = self.tensor_to_image(image)
+
+            temp_directory = tempfile.gettempdir()
+            unique_suffix = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5))
+            filename = f"image{unique_suffix}.png"
+            img_path = os.path.join(temp_directory, filename)
+
+            pil_image.save(img_path)
+
+            # img to base64
+            img = cv2.imread(img_path)
+            img_data = cv2.imencode('test.PNG', img)[1].tobytes()
+            base64_data = base64.b64encode(img_data)
+            img_base64 = str(base64_data, encoding='utf-8')
+
             header_img = {
                 'Authorization': f'Bearer {token}',
                 'Content-Type': 'application/json'
             }
 
-            data_img = {"model": "glm-4v", "messages": [{"role": "user",
-                                                         "content": [{"type": "text", "text": f"{prompt}"},
-                                                                     {"type": "image_url",
-                                                                      "image_url": {"url": f"{img_url}"}}]}]}
+            data_img = {"model": "glm-4v", "messages": [{"role": "user","content": [{"type": "text", "text": f"{prompt}"},{"type": "image_url","image_url": {"url": f"{img_base64}"}}]}]}
             data_img = json.dumps(data_img)
 
             response = requests.post(url=url, headers=header_img, data=data_img)
@@ -147,6 +170,9 @@ class ZhipuaiApi_Img:
             content_img = content_img.get("message", 'Default Value')
             content_img = content_img.get("content", 'Default Value')
             prompt = re.sub(r'\n+', '', content_img)
+
+            os.remove(img_path)
+
         return (prompt,)