change dependence

pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "swanlab" # 项目名称
-dynamic = ["version", "dependencies", "readme"] # 动态配置版本，依赖，可选依赖，readme
+dynamic = ["version", "dependencies", "optional-dependencies", "readme"] # 动态配置版本，依赖，可选依赖，readme
 description = "Python library for streamlined tracking and management of AI training processes." # 项目描述
 license = "Apache-2.0" # 项目许可证
 requires-python = ">=3.8" # python版本要求，我们只维护python3.8以上版本
@@ -56,6 +56,9 @@ pattern = ".*\"version\":\\s*\"(?P<version>[^\"]+)\""
 [tool.hatch.metadata.hooks.requirements_txt]
 filename = "requirements.txt"
 
+[tool.hatch.metadata.hooks.requirements_txt.optional-dependencies]
+media = ["requirements-media.txt"]  # pip install "swanlab[meida]"
+
 [tool.hatch.metadata.hooks.fancy-pypi-readme]  # 动态设置readme
 content-type = "text/markdown"
 fragments = [{ path = "README.md" }]
@@ -87,7 +90,7 @@ include = [
     "/test", # 包含一些测试脚本，确保测试成功
     "/README.md", # 包含readme
     "/requirements.txt", # 包含依赖
-    "/requirements-swan.txt", # 包含可选依赖
+    "/requirements-media.txt", # 包含可选依赖
 ]
 
 [tool.hatch.build.targets.wheel]

requirements-meida.txt

@@ -0,0 +1,4 @@
+soundfile
+pillow
+matplotlib
+numpy

requirements.txt

@@ -1,26 +1,9 @@
 # for swanlab
-swankit==0.1.0b8
-swanboard==0.1.2b4
-
-# data process
-soundfile
-pillow
-matplotlib
-numpy
-
-# web server
+swankit==0.1.0b9
+swanboard==0.1.2b5
 cos-python-sdk-v5
-fastapi>=0.110.1
-uvicorn>=0.14.0
 requests
 click
-
-
-# database
-ujson
 pyyaml
-
-
-# Information collection
 psutil
 pynvml

swanlab/data/formater.py

@@ -9,7 +9,7 @@
 """
 import os
 import re
-import ujson
+import json
 import yaml
 
 
@@ -59,7 +59,7 @@ def check_load_json_yaml(file_path: str, param_name):
         raise PermissionError(
             "No permission to read {}, please check if you have the permission.".format(param_name)
         )
-    load = ujson.load if path_suffix == "json" else yaml.safe_load
+    load = json.load if path_suffix == "json" else yaml.safe_load
     with open(file_path, "r") as f:
         # 读取配置文件的内容
         file_data = load(f)

swanlab/data/modules/audio/__init__.py

@@ -7,11 +7,16 @@
 @Description:
     音频模块
 """
-from swankit.core.data import MediaType
-from swankit.core import MediaBuffer, DataSuite as D
+from swankit.core import MediaBuffer, DataSuite as D, MediaType
 from typing import Union
-import soundfile as sf
-import numpy as np
+
+try:
+    # noinspection PyPackageRequirements
+    import soundfile as sf
+    # noinspection PyPackageRequirements
+    import numpy as np
+except ImportError:
+    sf, np = None, None
 
 
 class Audio(MediaType):
@@ -34,6 +39,11 @@ def __init__(
         caption: str
             Caption for the audio.
         """
+        if sf is None or np is None:
+            raise ImportError(
+                "soundfile and numpy are required for Audio class, "
+                'you can install them by `pip install "swanlab[media]"`'
+            )
         super().__init__()
         if isinstance(data_or_path, str):
             # 如果输入为路径字符串

swanlab/data/modules/image/__init__.py

@@ -1,11 +1,17 @@
-import numpy as np
-from PIL import Image as PILImage
-from matplotlib import pyplot as plt
-from swankit.core.data import MediaType
-from swankit.core import MediaBuffer, DataSuite as D
+from swankit.core import MediaBuffer, DataSuite as D, MediaType
 from typing import Union, Any
 from io import BytesIO
 
+try:
+    # noinspection PyPackageRequirements
+    import numpy as np
+    # noinspection PyPackageRequirements
+    from PIL import Image as PILImage
+    # noinspection PyPackageRequirements
+    from matplotlib import pyplot as plt
+except ImportError:
+    np, PILImage, plt = None, None, None
+
 
 def is_pytorch_tensor_typename(typename: str) -> bool:
     return typename.startswith("torch.") and ("Tensor" in typename or "Variable" in typename)
@@ -82,7 +88,11 @@ def __init__(
 
             If it is None, it means no scaling for the image.
         """
-
+        if PILImage is None or np is None or plt is None:
+            raise ImportError(
+                "pillow、numpy and matplotlib are required for Image class, "
+                'you can install them by `pip install "swanlab[media]"`'
+            )
         super().__init__()
         self.format = self.__convert_file_type(file_type)
         self.size = convert_size(size)

swanlab/data/modules/wrapper.py

@@ -7,8 +7,7 @@
 @Description:
     包装器
 """
-from swankit.core import ParseResult, ParseErrorInfo
-from swankit.core.data import MediaType
+from swankit.core import ParseResult, ParseErrorInfo, MediaType
 from typing import Union, List, Optional
 from swanlab.error import DataTypeError
 from .line import Line

swanlab/data/run/system/info.py

@@ -23,7 +23,7 @@ def __replace_second_colon(input_string, replacement):
     if first_colon_index != -1:
         second_colon_index = input_string.find(":", first_colon_index + 1)
         if second_colon_index != -1:
-            return input_string[:second_colon_index] + replacement + input_string[second_colon_index + 1 :]
+            return input_string[:second_colon_index] + replacement + input_string[second_colon_index + 1:]
     return input_string
 
 
@@ -108,7 +108,7 @@ def __get_nvidia_gpu_info():
                 gpu_name = gpu_name.decode("utf-8")
             info["type"].append(gpu_name)
             # 获取 GPU 的总显存, 单位为GB
-            info["memory"].append(round(pynvml.nvmlDeviceGetMemoryInfo(handle).total / (1024**3)))
+            info["memory"].append(round(pynvml.nvmlDeviceGetMemoryInfo(handle).total / (1024 ** 3)))
 
     except pynvml.NVMLError as e:
         swanlog.debug(f"An error occurred when getting GPU info: {e}")
@@ -120,17 +120,17 @@ def __get_nvidia_gpu_info():
 
 
 def __get_apple_gpu_info():
-    import ujson
+    import json
 
     info = {"cores": None, "type": [], "memory": []}
 
     # 使用system_profiler命令以JSON格式获取GPU信息
     try:
         result = subprocess.run(["system_profiler", "SPHardwareDataType", "-json"], capture_output=True, text=True)
-        gpu_name = ujson.loads(result.stdout)["SPHardwareDataType"][0]["chip_type"]
-        memory = ujson.loads(result.stdout)["SPHardwareDataType"][0]["physical_memory"]
+        gpu_name = json.loads(result.stdout)["SPHardwareDataType"][0]["chip_type"]
+        memory = json.loads(result.stdout)["SPHardwareDataType"][0]["physical_memory"]
         memory = str(memory).lower().replace("gb", "")
-        number_processors = ujson.loads(result.stdout)["SPHardwareDataType"][0]["number_processors"]
+        number_processors = json.loads(result.stdout)["SPHardwareDataType"][0]["number_processors"]
     except:
         return None
 
@@ -147,7 +147,7 @@ def __get_apple_gpu_info():
     # try:
     #     command = [str(binary_path), "--json"]
     #     output = (subprocess.check_output(command, universal_newlines=True).strip().split("\n"))[0]
-    #     raw_stats = ujson.loads(output)
+    #     raw_stats = json.loads(output)
     #     stats = {
     #         "gpu": raw_stats["utilization"],
     #         "memoryAllocated": raw_stats["mem_used"],
@@ -190,7 +190,7 @@ def __get_memory_size():
     try:
         # 获取系统总内存大小
         mem = psutil.virtual_memory()
-        total_memory = round(mem.total / (1024**3))  # 单位为GB
+        total_memory = round(mem.total / (1024 ** 3))  # 单位为GB
         return total_memory
     except Exception as e:
         swanlog.debug(f"An error occurred when getting memory size: {e}")