OneDiff加速AnimateDiff-CLI-Prompt-Travel的技术实践

背景介绍

AnimateDiff-CLI-Prompt-Travel是一个基于Diffusers框架实现的动画生成工具，它通过自定义的UNet3DConditionModel实现了视频序列的生成。在实际应用中，我们发现其推理速度有待提升，因此尝试使用OneDiff的oneflow_compile功能对其进行加速优化。

技术挑战与解决方案

1. 自定义UNet模型的编译问题

原项目使用了自定义实现的UNet3DConditionModel，而非Diffusers官方实现。直接使用oneflow_compile会遇到算子不支持的问题。

解决方案：

• 确认OneDiff版本：从dev_comfyui_animatediff_evolved分支切换到master分支
• 确保编译环境配置正确

2. 最临近插值算子不支持问题

原代码中使用了torch.nn.functional.interpolate的最临近插值功能，这是OneDiff当前不支持的算子。

解决方案： 我们实现了自定义的最临近插值函数：

def nearest_interpolate(input_tensor, size=None, scale_factor=None):
    """
    最临近插值实现
    :param input_tensor: 输入张量
    :param size: 目标尺寸
    :param scale_factor: 缩放因子
    :return: 插值后的张量
    """
    new_shape = list(input_tensor.shape)
    if size is not None:
        for i in range(2, 2 + len(size)):
            new_shape[i] = size[i - 2]
    if scale_factor is not None:
        for i in range(2, 2 + len(scale_factor)):
            new_shape[i] = int(new_shape[i] * scale_factor[i - 2])
    
    output_tensor = torch.zeros(new_shape, dtype=input_tensor.dtype, device=input_tensor.device.type)
    index_k = (torch.arange(0, output_tensor.shape[-1]) / output_tensor.shape[-1] * input_tensor.shape[-1]).int()
    index_j = (torch.arange(0, output_tensor.shape[-2]) / output_tensor.shape[-2] * input_tensor.shape[-2]).int()
    index_i = (torch.arange(0, output_tensor.shape[-3]) / output_tensor.shape[-3] * input_tensor.shape[-3]).int()
    
    output_tensor = input_tensor[:, :, index_i][:, :, :, index_j][:, :, :, :, index_k]
    return output_tensor

实现要点：

1. 避免使用torch.tensor()类型转换，改用.int()方法
2. 采用索引方式实现插值，而非直接使用numpy转换
3. 支持3D张量的插值操作

3. 其他优化技巧

• 同时对UNet和ControlNet部分进行编译优化
• 注意数据类型和设备的一致性
• 监控显存使用情况，避免因优化导致显存溢出

性能评估

经过优化后：

1. 模型推理速度得到显著提升
2. 自定义插值函数比原生实现慢约5倍，但整体仍带来正向收益
3. 系统稳定性良好，无崩溃现象

总结与展望

通过OneDiff的编译优化，我们成功提升了AnimateDiff-CLI-Prompt-Travel的推理性能。未来可考虑：

1. 进一步优化自定义算子的实现效率
2. 探索更多可编译的模型组件
3. 研究混合精度训练带来的额外性能提升

这种技术方案不仅适用于本项目，也可推广到其他基于Diffusers的自定义模型优化场景中。