cv_unet_image-colorization部署实录：CentOS 7 + CUDA 11.8兼容性适配全过程

1. 项目概述与核心价值

cv_unet_image-colorization是一个基于深度学习的黑白照片上色工具，能够将老旧的黑白照片自动转换为彩色图像。这个工具特别适合处理历史照片、家庭老照片等珍贵影像资料，让褪色的记忆重新焕发光彩。

在实际部署过程中，我们遇到了PyTorch版本兼容性问题。新版本的PyTorch（2.6+）在加载旧模型时会出现安全限制报错，这需要特殊的兼容性处理。本文记录了从环境准备到最终部署的完整过程，重点解决了这些技术难题。

工具的核心优势包括：

• 完全本地运行：所有处理都在本地完成，无需上传到云端，保护用户隐私
• GPU加速：利用CUDA和GPU进行加速处理，大幅提升上色速度
• 简单易用：通过Streamlit提供直观的网页界面，无需技术背景也能轻松使用
• 高质量上色：基于先进的生成对抗网络，能够产生自然逼真的上色效果

2. 环境准备与系统要求

2.1 硬件要求

要顺利运行这个上色工具，你的系统需要满足以下硬件要求：

• 操作系统：CentOS 7.x（其他Linux发行版也可参考）
• GPU：NVIDIA显卡，至少4GB显存（GTX 1060以上推荐）
• 内存：8GB以上（处理大图片时需要更多内存）
• 存储空间：至少10GB可用空间（用于存放模型和依赖包）

2.2 软件依赖

在开始安装前，请确保系统已安装以下基础软件：

# 更新系统包
sudo yum update -y

# 安装基础开发工具
sudo yum groupinstall "Development Tools" -y
sudo yum install openssl-devel libffi-devel bzip2-devel -y

# 安装Python 3.8+
sudo yum install python3 python3-devel -y

3. CUDA与PyTorch环境配置

3.1 CUDA 11.8安装

CentOS 7上安装CUDA 11.8需要一些额外的步骤：

# 首先安装EPEL仓库
sudo yum install epel-release -y

# 安装NVIDIA驱动依赖
sudo yum install kernel-devel kernel-headers gcc make dkms -y

# 下载CUDA 11.8安装包
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run

# 给安装文件添加执行权限
sudo chmod +x cuda_11.8.0_520.61.05_linux.run

# 运行安装程序
sudo ./cuda_11.8.0_520.61.05_linux.run

安装过程中，记得选择安装驱动程序和CUDA工具包。安装完成后，需要设置环境变量：

# 编辑bash配置文件
echo 'export PATH=/usr/local/cuda-11.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc

# 使配置生效
source ~/.bashrc

3.2 PyTorch与依赖安装

由于兼容性要求，我们需要安装特定版本的PyTorch：

# 创建Python虚拟环境
python3 -m venv colorization_env
source colorization_env/bin/activate

# 安装PyTorch与CUDA 11.8兼容版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装其他依赖
pip install modelscope==1.10.0 streamlit==1.29.0 opencv-python==4.8.1.78 pillow==10.0.1

4. 核心兼容性问题解决

4.1 PyTorch 2.6+兼容性修复

新版本PyTorch引入了安全限制，导致无法直接加载旧模型。我们需要重写torch.load方法：

import torch
import warnings

# 保存原始load方法
_original_torch_load = torch.load

# 定义兼容性load函数
def compatible_load(f, map_location=None, pickle_module=None, **kwargs):
    # 忽略weights_only参数，强制使用旧式加载
    if 'weights_only' in kwargs:
        kwargs.pop('weights_only')
        warnings.warn("忽略weights_only参数以兼容旧模型", UserWarning)
    
    return _original_torch_load(f, map_location, pickle_module, **kwargs)

# 替换torch.load方法
torch.load = compatible_load

4.2 模型加载适配代码

在实际使用中，我们需要确保模型正确加载：

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

def load_colorization_model():
    """加载上色模型并处理兼容性问题"""
    try:
        # 应用兼容性修复
        apply_compatibility_fix()
        
        # 创建上色pipeline
        colorizer = pipeline(
            Tasks.image_colorization,
            model='damo/cv_unet_image-colorization',
            device='gpu'  # 强制使用GPU
        )
        return colorizer
    except Exception as e:
        print(f"模型加载失败: {str(e)}")
        return None

def apply_compatibility_fix():
    """应用PyTorch兼容性修复"""
    import torch
    import functools
    
    # 备份原始load方法
    original_load = torch.load
    
    # 定义修复后的load方法
    @functools.wraps(original_load)
    def fixed_load(*args, **kwargs):
        if 'weights_only' in kwargs:
            kwargs['weights_only'] = False
        return original_load(*args, **kwargs)
    
    # 替换torch.load
    torch.load = fixed_load

5. Streamlit界面开发与优化

5.1 基础界面布局

创建一个直观易用的网页界面：

import streamlit as st
import tempfile
from PIL import Image
import os

def main():
    st.set_page_config(
        page_title="黑白照片上色工具",
        page_icon="🎨",
        layout="wide"
    )
    
    st.title("🎨 黑白照片上色工具")
    st.markdown("上传黑白或老照片，AI自动为其添加逼真色彩")
    
    # 创建侧边栏
    with st.sidebar:
        st.header("上传照片")
        uploaded_file = st.file_uploader(
            "选择一张黑白/老照片",
            type=['jpg', 'jpeg', 'png'],
            help="支持JPG、PNG格式，最大20MB"
        )
        
        if uploaded_file is not None:
            st.success("照片上传成功！")
            if st.button("开始上色", type="primary"):
                process_image(uploaded_file)
    
    # 主界面布局
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("原始照片")
        if uploaded_file is not None:
            st.image(uploaded_file, use_column_width=True)
        else:
            st.info("请在左侧上传黑白照片")
    
    with col2:
        st.subheader("上色结果")
        if 'result_image' in st.session_state:
            st.image(st.session_state.result_image, use_column_width=True)
            st.success("处理完成！")
        else:
            st.info("上色结果将在这里显示")

if __name__ == "__main__":
    main()

5.2 图像处理逻辑

实现完整的图像处理流程：

def process_image(uploaded_file):
    """处理上传的图像"""
    try:
        # 显示处理中状态
        with st.spinner('AI正在为照片上色，请稍候...'):
            # 保存上传的文件
            with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file:
                tmp_file.write(uploaded_file.getvalue())
                temp_path = tmp_file.name
            
            # 读取图像
            image = Image.open(temp_path)
            
            # 确保模型已加载
            if 'colorizer' not in st.session_state:
                st.session_state.colorizer = load_colorization_model()
            
            # 执行上色
            result = st.session_state.colorizer(image)
            
            # 保存结果到session state
            st.session_state.result_image = result['output_img']
            
            # 清理临时文件
            os.unlink(temp_path)
            
    except Exception as e:
        st.error(f"处理失败: {str(e)}")

6. 完整部署脚本

创建一个一键部署和启动的脚本：

#!/bin/bash
# deploy_colorization.sh

echo "开始部署黑白照片上色工具..."

# 检查CUDA是否可用
if ! command -v nvcc &> /dev/null; then
    echo "错误: 未检测到CUDA，请先安装CUDA 11.8"
    exit 1
fi

echo "CUDA版本: $(nvcc --version | grep release)"

# 创建项目目录
PROJECT_DIR="$HOME/image-colorization"
mkdir -p $PROJECT_DIR
cd $PROJECT_DIR

# 创建虚拟环境
echo "创建Python虚拟环境..."
python3 -m venv venv
source venv/bin/activate

# 安装依赖
echo "安装依赖包..."
pip install --upgrade pip
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install modelscope==1.10.0 streamlit==1.29.0 opencv-python==4.8.1.78 pillow==10.0.1

# 下载示例代码
echo "下载应用代码..."
cat > app.py << 'EOF'
# 这里插入前面完整的Python代码
EOF

echo "部署完成！"
echo "启动命令: source venv/bin/activate && streamlit run app.py"

7. 常见问题与解决方案

7.1 CUDA相关错误

问题：CUDA out of memory

• 原因：显存不足，图片太大或模型太大
• 解决：减小输入图片尺寸，使用更小的模型版本

# 在处理前调整图片大小
def resize_image(image, max_size=1024):
    """调整图片大小以避免显存不足"""
    width, height = image.size
    if max(width, height) > max_size:
        scale = max_size / max(width, height)
        new_width = int(width * scale)
        new_height = int(height * scale)
        return image.resize((new_width, new_height), Image.LANCZOS)
    return image

7.2 模型加载失败

问题：模型下载失败或加载超时

• 原因：网络问题或模型服务器问题
• 解决：使用本地模型或设置超时重试

# 设置模型下载超时和重试
import requests
from modelscope.hub.snapshot_download import snapshot_download

# 设置超时和重试
os.environ['MODELSCOPE_REQUEST_TIMEOUT'] = '300'
os.environ['MODELSCOPE_DOWNLOAD_TIMEOUT'] = '300'

# 或者预先下载模型
model_dir = snapshot_download('damo/cv_unet_image-colorization')

8. 总结与使用建议

通过本文的步骤，你应该已经成功在CentOS 7上部署了cv_unet_image-colorization工具。这个工具不仅解决了PyTorch新版本的兼容性问题，还提供了简单易用的网页界面。

使用建议：

1. 图片选择：选择清晰度高、对比度好的黑白照片，效果更佳
2. 处理时间：首次运行需要下载模型，后续处理通常需要10-30秒
3. 结果保存：处理完成后可以右键保存上色结果
4. 批量处理：如果需要处理大量照片，可以考虑编写批量处理脚本

性能优化提示：

• 如果处理速度较慢，可以尝试减小输入图片的尺寸
• 确保GPU驱动和CUDA版本正确安装
• 关闭其他占用GPU的应用程序以释放显存

这个工具特别适合：

• 家庭老照片修复和数字化
• 历史档案资料的颜色还原
• 艺术创作中的黑白照片上色
• 教育和研究用途

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。