pycoze 0.1.75__py3-none-any.whl → 0.1.76__py3-none-any.whl

pycoze/ai/__init__.py

@@ -1 +1 @@
-from .gpu_reserve import reserve_gpu_retry, release_gpu, reserve_gpu
+from .vram_reserve import reserve_vram_retry, release_vram, reserve_vram

pycoze/ai/comfyui.py

@@ -0,0 +1,33 @@
+from .vram_reserve import reserve_vram_retry, release_vram
+
+
+class ComfyUI:
+    def __init__(self, file_path):
+        self.file_path = file_path
+
+    def load(self):
+        reserve_vram_retry(20)  # 20GB VRAM
+        with open(self.file_path, "r") as file:
+            return file.read()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        clear_memory()
+        return None
+
+    def unload(self):
+        clear_memory()
+        return None
+
+
+def clear_memory():
+    import requests
+    import json
+
+    release_vram()
+
+    url = "http://127.0.0.1:8188/free"
+    data = {"unload_models": True, "free_memory": False}
+    response = requests.post(
+        url, data=json.dumps(data), headers={"Content-Type": "application/json"}
+    )
+    return response.status_code, response.text

pycoze/ai/{gpu_reserve.py → vram_reserve.py}

@@ -12,14 +12,14 @@ try:
     # 定义数据库连接和初始化
     params = utils.arg.read_params_file()
     if params:
-        DATABASE_PATH = params["appPath"] + "/gpu_usage.db"
+        DATABASE_PATH = params["appPath"] + "/vram_usage.db"
     else:
         raise Exception("No params")
 except:
     DATABASE_DIR = os.path.expanduser("~/pycoze")
     os.makedirs(DATABASE_DIR, exist_ok=True)
-    DATABASE_PATH = os.path.join(DATABASE_DIR, "gpu_usage.db")
-TABLE_NAME = "gpu_usage"
+    DATABASE_PATH = os.path.join(DATABASE_DIR, "vram_usage.db")
+TABLE_NAME = "vram_usage"
 
 
 def initialize_db():
@@ -38,10 +38,9 @@ def initialize_db():
     conn.close()
 
 
-# 检测GPU资源
-def get_gpu_resources():
+def get_vram_resources():
     try:
-        # 使用nvidia-smi命令获取GPU信息
+        # 使用nvidia-smi命令获取VRAM信息
         result = subprocess.run(
             ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
             stdout=subprocess.PIPE,
@@ -55,7 +54,7 @@ def get_gpu_resources():
             except:
                 pass
 
-        # 获取正在使用GPU的进程信息
+        # 获取正在使用VRAM的进程信息
         process_result = subprocess.run(
             [
                 "nvidia-smi",
@@ -76,18 +75,17 @@ def get_gpu_resources():
                     python_memory_usage += float(used_memory)
                 except:
                     pass
-        print("total_free_gpu_memory: ", total_free_memory)
-        print("python_gpu_memory_usage: ", python_memory_usage)
+        print("total_free_vram_memory: ", total_free_memory)
+        print("python_vram_memory_usage: ", python_memory_usage)
         # 计算排除python进程后的总空闲内存
         total_free_memory -= python_memory_usage
         return round(total_free_memory / 1024, 2)
     except Exception as e:
-        print(f"Error getting GPU resources: {e}")
+        print(f"Error getting VRAM resources: {e}")
         return 0.0
 
 
-# 预留GPU资源
-def reserve_gpu(gb, uid=None):
+def reserve_vram(gb, uid=None):
     if uid is None:
         uid = f"pid:{os.getpid()}"
     with sqlite3.connect(DATABASE_PATH) as conn:
@@ -96,21 +94,21 @@ def reserve_gpu(gb, uid=None):
         total_reserved = cursor.fetchone()[0]
         if total_reserved is None:
             total_reserved = 0.0
-        available_gb = get_gpu_resources() - total_reserved
+        available_gb = get_vram_resources() - total_reserved
         if available_gb >= gb:
             cursor.execute(
                 f"INSERT INTO {TABLE_NAME} (uid, reserved_gb) VALUES (?, ?)",
                 (uid, gb),
             )
             conn.commit()
-            print(f"预留成功，剩余GPU大小: {available_gb - gb} GB")
+            print(f"预留成功，剩余VRAM大小: {available_gb - gb} GB")
             return True
         else:
-            print(f"预留失败，剩余GPU大小: {available_gb} GB")
+            print(f"预留失败，剩余VRAM大小: {available_gb} GB")
             return False
 
 
-def reserve_gpu_retry(gb, retry=None, uid=None):
+def reserve_vram_retry(gb, retry=None, uid=None):
     if retry is None:
         # 接近无限重试，python中允许无限大的整数，尽管sys.maxsize不是真正的无限大，但足够大
         retry = sys.maxsize
@@ -118,32 +116,31 @@ def reserve_gpu_retry(gb, retry=None, uid=None):
         time.sleep(1)
         if i % 10 == 0 or i < 10:
             print(f"重试第{i}次")
-        if reserve_gpu(gb, uid):
+        if reserve_vram(gb, uid):
             return True
     return False
 
 
-# 释放GPU资源
-def release_gpu(uid=None):
+def release_vram(uid=None):
     if uid is None:
         uid = f"pid:{os.getpid()}"
     with sqlite3.connect(DATABASE_PATH) as conn:
         cursor = conn.cursor()
         cursor.execute(f"DELETE FROM {TABLE_NAME} WHERE uid = ?", (uid,))
         conn.commit()
-        # 计算释放后的剩余GPU大小
+        # 计算释放后的剩余VRAM大小
         cursor.execute(f"SELECT SUM(reserved_gb) FROM {TABLE_NAME}")
         total_reserved = cursor.fetchone()[0]
         if total_reserved is None:
             total_reserved = 0.0
-        available_gb = get_gpu_resources() - total_reserved
-        print(f"释放成功，剩余GPU大小: {available_gb} GB")
+        available_gb = get_vram_resources() - total_reserved
+        print(f"释放成功，剩余VRAM大小: {available_gb} GB")
 
 
 # 注册退出时的清理函数
 def cleanup():
-    release_gpu()
-    print("程序退出，GPU资源已释放")
+    release_vram()
+    print("程序退出，VRAM资源已释放")
 
 
 def initialize_and_check():
@@ -173,12 +170,12 @@ initialize_and_check()
 atexit.register(cleanup)
 
 if __name__ == "__main__":
-    if reserve_gpu_retry(5):
-        print("(1)GPU资源预留成功")
-        if reserve_gpu_retry(5):
-            print("(2)GPU资源预留成功")
+    if reserve_vram_retry(5):
+        print("(1)VRAM资源预留成功")
+        if reserve_vram_retry(5):
+            print("(2)VRAM资源预留成功")
             time.sleep(100)
-            release_gpu()
-            print("GPU资源释放成功")
+            release_vram()
+            print("VRAM资源释放成功")
     else:
-        print("GPU资源不足，无法预留")
+        print("VRAM资源不足，无法预留")

pycoze/bot/bot.py

@@ -36,7 +36,7 @@ def agent_chat(bot_setting_file, history):
             model=cfg["model"],
             temperature=(
                 role_setting["temperature"] * 2
-                if cfg["model"].startswith("deepseek")
+                if cfg["model"].startswith("deepseek") or cfg["model"].startswith("yi-")
                 else role_setting["temperature"]
             ),
             stop_sequences=[
@@ -45,7 +45,11 @@ def agent_chat(bot_setting_file, history):
             ],  # 停用deepseek的工具调用标记，不然会虚构工具调用过程和结果
         )
     prompt = role_setting["prompt"]
-    if cfg["model"].startswith("deepseek") and len(tools) > 0:
+    if (
+        cfg["model"].startswith("deepseek")
+        or cfg["model"].startswith("yi-")
+        and len(tools) > 0
+    ):
         prompt += """
 如果需要调用工具，请使用以正确的json格式进行结尾（务必保证json格式正确，不要出现反斜杠未转义等问题）：
 ```json

{pycoze-0.1.75.dist-info → pycoze-0.1.76.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycoze
-Version: 0.1.75
+Version: 0.1.76
 Summary: Package for pycoze only!
 Author: Yuan Jie Xiong
 Author-email: aiqqqqqqq@qq.com

{pycoze-0.1.75.dist-info → pycoze-0.1.76.dist-info}/RECORD

@@ -2,10 +2,11 @@ pycoze/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycoze/module.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycoze/access/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycoze/access/tool_for_bot.py,sha256=Nv6XlxcqaLaMVQNDkpcthlqaSxyfHoOfWiTScQXO22Y,2656
-pycoze/ai/__init__.py,sha256=wUiH-Oy_YbaoiqJN2f9ZvFuRCv9oEDP-oNs5Q8y9IzA,69
-pycoze/ai/gpu_reserve.py,sha256=WcD3-K9TsapSXCNVUZFm0k-GHZxznWdFBcljA41dMp8,5764
+pycoze/ai/__init__.py,sha256=Smivpb8qbRnzWkzKRe2IxsmKP5Dh8EvngDFdkD_DVLo,73
+pycoze/ai/comfyui.py,sha256=u75tZywkuXiOdm7XST2kBAaveJKpPvY_qTQr_TN9sXk,795
+pycoze/ai/vram_reserve.py,sha256=s55Cy-Q5mTq-k5oIPbAFwCfrjatjN0QTjQxW7WBTPZI,5738
 pycoze/bot/__init__.py,sha256=6HHMxDQVOyZM9dtSjQm9tjGnhj4h7CixD0JOvEwTi48,41
-pycoze/bot/bot.py,sha256=iJ60C5YOiiAwADQAdWlD15uAgPUuSap5O7gaiJaXQ-U,2769
+pycoze/bot/bot.py,sha256=aEG_B4ML6ONZ6FGcWYfBtlcMSuYBE8CjXo6ZYESQEaU,2869
 pycoze/bot/agent/__init__.py,sha256=YR9vpkEQn1e4937r_xFPJXUCPBEJ0SFzEQDBe2x3-YA,157
 pycoze/bot/agent/agent.py,sha256=7xnz4a8LDyAVU0O3tS6n7_jSnr_lukQxfNR2znCWoV4,4611
 pycoze/bot/agent/assistant.py,sha256=QLeWaPi415P9jruYOm8qcIbC94cXXAhJYmLTkyC9NTQ,1267
@@ -22,8 +23,8 @@ pycoze/ui/ui_def.py,sha256=UhhU_yB3GV9ISbvTWT48hsHPHI250BhMILh6bu5Uioo,4206
 pycoze/utils/__init__.py,sha256=TNJhFfY7JYdLlzuP9GvgxfNXUtbgH_NUUJSqHXCxJn4,78
 pycoze/utils/arg.py,sha256=kA3KBQzXc2WlH5XbF8kfikfpqljiKaW7oY_GE4Qyffc,753
 pycoze/utils/text_or_file.py,sha256=gpxZVWt2DW6YiEg_MnMuwg36VNf3TX383QD_1oZNB0Y,551
-pycoze-0.1.75.dist-info/LICENSE,sha256=QStd_Qsd0-kAam_-sOesCIp_uKrGWeoKwt9M49NVkNU,1090
-pycoze-0.1.75.dist-info/METADATA,sha256=whsTD8FRG8nXKTOU3RrlInnqBx73vvRt21FNcMJ4qOw,719
-pycoze-0.1.75.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-pycoze-0.1.75.dist-info/top_level.txt,sha256=76dPeDhKvOCleL3ZC5gl1-y4vdS1tT_U1hxWVAn7sFo,7
-pycoze-0.1.75.dist-info/RECORD,,
+pycoze-0.1.76.dist-info/LICENSE,sha256=QStd_Qsd0-kAam_-sOesCIp_uKrGWeoKwt9M49NVkNU,1090
+pycoze-0.1.76.dist-info/METADATA,sha256=ZPC58zR5Sued7MLMyaiOPY7T7PYzd-MCadxAStH_HV0,719
+pycoze-0.1.76.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+pycoze-0.1.76.dist-info/top_level.txt,sha256=76dPeDhKvOCleL3ZC5gl1-y4vdS1tT_U1hxWVAn7sFo,7
+pycoze-0.1.76.dist-info/RECORD,,