ai-swap/backend/modules/comfyui/image_tools.py

import io
import uuid
import requests
import numpy as np
from PIL import Image
from backend.utils.logger_config import setup_logger
from backend.utils.system_config import Config

logger = setup_logger(__name__)

def show_gif(fname: str):
    """
    显示图片（兼容非IPython环境）
    Args:
        fname: 图片文件路径
    """
    try:
        import base64
        from IPython import display
        with open(fname, 'rb') as fd:
            b64 = base64.b64encode(fd.read()).decode('ascii')
        return display.HTML(f'<img src="data:image/gif;base64,{b64}" />')
    except ImportError:
        # 如果没有IPython，只记录日志，不显示图片
        logger.info(f"图片已保存: {fname} (IPython不可用，跳过显示)")
        return None
    except Exception as e:
        logger.warning(f"显示图片失败: {e}")
        return None

def upload_image(image_1: np.ndarray, config: Config) -> str:
    """
    将图片上传到服务器
    Args:
        image_1: 图片
        config: 配置
    Returns:
        image_name: 图片名称
    """
    image = Image.fromarray(image_1)
    buffer = io.BytesIO()
    image.save(buffer, format="png")
    image = buffer.getbuffer()
    image_hash = hash(image.tobytes())
    
    if image_hash in config.uploaded_image:
        return config.uploaded_image[image_hash]
    
    image_name = str(uuid.uuid4()) + ".png"
    config.uploaded_image[image_hash] = image_name
    image_file = {"image": (image_name, image)}
    
    try:
        requests.post(url=f"http://{config.server_address}/upload/image", files=image_file)
    except requests.RequestException as e:
        logger.error(f"Image upload failed: {e}")
        raise
    
    return image_name


def save_uploaded_file(uploaded_file):
    """保存上传的文件并返回图片对象"""
    if uploaded_file is not None:
        try:
            image_bytes = uploaded_file.getvalue()
            image = Image.open(io.BytesIO(image_bytes))
            if image.mode == 'RGBA':
                image = image.convert('RGB')
            return image
        except Exception as e:
            return None
    return None

def image_to_bytes(image):
    """将PIL Image对象转换为字节"""
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    return img_byte_arr.getvalue()

def bytes_to_image(bytes_data):
    """将字节转换为PIL Image对象"""
    return Image.open(io.BytesIO(bytes_data))