Upload utils.py

utils.py

@@ -0,0 +1,194 @@
+import re
+import json
+import requests
+from io import BytesIO
+from PIL import Image, ImageOps, ImageSequence
+import numpy as np
+import torch
+import base64
+from PIL import ImageFile, UnidentifiedImageError
+import hashlib
+import time
+
+
+def pillow(fn, arg):
+    prev_value = None
+    try:
+        x = fn(arg)
+    except (OSError, UnidentifiedImageError, ValueError): #PIL issues #4472 and #2445, also fixes ComfyUI issue #3416
+        prev_value = ImageFile.LOAD_TRUNCATED_IMAGES
+        ImageFile.LOAD_TRUNCATED_IMAGES = True
+        x = fn(arg)
+    finally:
+        if prev_value is not None:
+            ImageFile.LOAD_TRUNCATED_IMAGES = prev_value
+    return x
+
+
+def image_file_to_base64(path: str) -> str:
+    """
+    读取本地图片文件，并返回其 Base64 编码字符串（不带 data URI 头）。
+    """
+    with open(path, 'rb') as f:
+        data = f.read()
+    # 将二进制数据编码为 Base64 字符串，并解码为 str
+    base64_str = base64.b64encode(data).decode('utf-8')
+    return base64_str
+
+
+def tensor_to_url(img_tensor: torch.Tensor) -> str:
+    # 批量→单张
+    if img_tensor.dim() == 4:
+        img_tensor = img_tensor[0]
+
+    # 确保在 CPU、float
+    img = img_tensor.detach().cpu().float()
+
+    # 只有三维时才考虑 permute
+    if img.dim() == 3:
+        d0, d1, d2 = img.shape
+        # 如果第一维是通道（常见 1,3,4），就把 (C,H,W)→(H,W,C)
+        if d0 in (1, 3, 4):
+            img = img.permute(1, 2, 0)
+        # 否则，如果最后一维是通道，就假设已经是 (H,W,C) 了
+        elif d2 in (1, 3, 4):
+            pass
+        else:
+            raise ValueError(f"Unexpected tensor shape {img.shape} for image conversion.")
+
+    # 此时 img.shape 应该是 (H, W) 或 (H, W, C)
+    arr = img.clamp(0, 1).mul(255).byte().numpy()
+
+    # 创建 PIL 图像
+    pil_img = Image.fromarray(arr)
+
+    buf = BytesIO()
+    pil_img.save(buf, format="JPEG")   # 或 "JPEG"
+    buf.seek(0)
+
+    # 3. 用 requests 传文件
+    files = {
+        "file": ("test.png", buf, "image/jpeg")  # (文件名, 文件内容, MIME类型)
+    }
+
+    UPLOAD_URL = "http://deepnet.meitustat.com/data-center/v1/cloud/upload"
+    
+    # 超时重试配置
+    max_retries = 3
+    timeout = 30  # 30秒超时
+    retry_delay = 1  # 重试间隔1秒
+    
+    for attempt in range(max_retries):
+        try:
+            response = requests.post(UPLOAD_URL, files=files, timeout=timeout)
+            response.raise_for_status()  # Raise exception for HTTP errors
+            return response.json()['data']['url']
+        except Exception as e:
+            if attempt == max_retries - 1:  # 最后一次尝试失败
+                raise Exception(f"上传失败，已重试 {max_retries} 次: {str(e)}")
+            print(f"上传失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}，{retry_delay}秒后重试...")
+            time.sleep(retry_delay)
+            retry_delay *= 2  # 指数退避，下次重试间隔翻倍
+
+
+def tensor_to_base64(img_tensor: torch.Tensor) -> str:
+    # 批量→单张
+    if img_tensor.dim() == 4:
+        img_tensor = img_tensor[0]
+
+    # 确保在 CPU、float
+    img = img_tensor.detach().cpu().float()
+
+    # 只有三维时才考虑 permute
+    if img.dim() == 3:
+        d0, d1, d2 = img.shape
+        # 如果第一维是通道（常见 1,3,4），就把 (C,H,W)→(H,W,C)
+        if d0 in (1, 3, 4):
+            img = img.permute(1, 2, 0)
+        # 否则，如果最后一维是通道，就假设已经是 (H,W,C) 了
+        elif d2 in (1, 3, 4):
+            pass
+        else:
+            raise ValueError(f"Unexpected tensor shape {img.shape} for image conversion.")
+
+    # 此时 img.shape 应该是 (H, W) 或 (H, W, C)
+    arr = img.clamp(0, 1).mul(255).byte().numpy()
+
+    # 创建 PIL 图像
+    pil_img = Image.fromarray(arr)
+
+    # 存到内存、再转 Base64
+    buf = BytesIO()
+    pil_img.save(buf, format="JPEG")
+    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+    return b64
+
+
+def load_image_from_url(image_url):
+    # 下载图片数据
+    response = requests.get(image_url)
+    response.raise_for_status()
+    img_data = BytesIO(response.content)
+
+    # 使用 PIL 打开图像
+    img = pillow(Image.open, img_data)
+
+    output_images = []
+    output_masks = []
+    w, h = None, None
+
+    excluded_formats = ['MPO']
+
+    for i in ImageSequence.Iterator(img):
+        i = pillow(ImageOps.exif_transpose, i)
+
+        if i.mode == 'I':
+            i = i.point(lambda i: i * (1 / 255))
+        image = i.convert("RGB")
+
+        if len(output_images) == 0:
+            w, h = image.size
+
+        if image.size != (w, h):
+            continue
+
+        image_np = np.array(image).astype(np.float32) / 255.0
+        image_tensor = torch.from_numpy(image_np)[None, ...]
+
+        if 'A' in i.getbands():
+            mask_np = np.array(i.getchannel('A')).astype(np.float32) / 255.0
+            mask_tensor = 1. - torch.from_numpy(mask_np)
+        else:
+            mask_tensor = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
+
+        output_images.append(image_tensor)
+        output_masks.append(mask_tensor.unsqueeze(0))
+
+    if len(output_images) > 1 and getattr(img, "format", None) not in excluded_formats:
+        output_image = torch.cat(output_images, dim=0)
+        output_mask = torch.cat(output_masks, dim=0)
+    else:
+        output_image = output_images[0]
+        output_mask = output_masks[0]
+
+    return output_image
+
+
+def extract_json_from_text(text: str) -> dict:
+    """
+    从给定的文本中提取第一个 JSON 对象并解析成 Python 字典。
+    
+    :param text: 包含 JSON 的文本
+    :return: 解析后的字典，如果未找到 JSON 则返回空字典
+    """
+    # 匹配第一个花括号及其中所有内容
+    match = re.search(r"\{[\s\S]*?\}", text)
+    if not match:
+        return {}
+    
+    json_str = match.group(0)
+    try:
+        return json.loads(json_str)
+    except json.JSONDecodeError as e:
+        print(f"JSON 解析失败: {e}")
+        return {}