susie复现流程

论文：Zero-Shot Robotic Manipulation With Pretrained Image-Editing Diffusion Models

代码地址：https://github.com/kvablack/susie

复现流程

1.创建环境

conda create -n susie python=3.10 -y
conda activate susie

2.安装Jax,选择的是Ubuntu系统安装的，显卡是3080，太低可能不行

pip install -U "jax[cuda12]"

3.切换到susie所在的文件夹，然后pip install -r requirements.txt, 并安装susie

(susie) zlm@zlm-MS-7C82:/home/new1/susie$ 
pip install -r requirements.txt 
pip install -e .

4.创建python文件运行示例代码

touch generate_image.py
python generate_image.py

并复制示例代码到python文件中，我修改了代码的图片显示部分，原本的代码是在Jupyter notebook中显示的。
from susie.model import create_sample_fn
from susie.jax_utils import initialize_compilation_cache
import requests
import numpy as np
from PIL import Image

initialize_compilation_cache()

IMAGE_URL = “https://rail.eecs.berkeley.edu/datasets/bridge_release/raw/bridge_data_v2/datacol2_toykitchen7/drawer_pnp/01/2023-04-19_09-18-15/raw/traj_group0/traj0/images0/im_12.jpg”

sample_fn = create_sample_fn(“kvablack/susie”)
image = np.array(Image.open(requests.get(IMAGE_URL, stream=True).raw).resize((256, 256)))
image_out = sample_fn(image, “open the drawer”)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5)) # 创建并排的两个子图
ax1.imshow(image)
ax1.set_title(‘Input Image’)
ax1.axis(‘off’) # 关闭坐标轴
ax2.imshow(image_out)
ax2.set_title(‘Output Image’)
ax2.axis(‘off’)

plt.tight_layout() # 自动调整子图间距
plt.show()

5. 出现的错误及对应解决办法

RuntimeError: jaxlib version 0.4.38 is newer than and incompatible with jax version 0.4.11. Please update your jax and/or jaxlib packages.

pip install --upgrade jax jaxlib

A module that was compiled using NumPy 1.x cannot be run in NumPy 2.2.5 as it may crash. To support both 1.x and 2.x versions of NumPy, modules must be compiled with NumPy 2.0. Some module may need to rebuild instead e.g. with ‘pybind11>=2.12’.

pip install "numpy<2"

ImportError: cannot import name ‘cached_download’ from ‘huggingface_hub’ (/home/zlm/anaconda3/envs/susie/lib/python3.10/site-packages/huggingface_hub/init.py),试了好几次，刚开始以为是huggingface_hub版本的原因，其实是diffusers。

pip install --upgrade diffusers

AttributeError: ‘Config’ object has no attribute ‘define_bool_state’

pip install --upgrade flax diffusers

AttributeError: module ‘jax.core’ has no attribute ‘Shape’

pip install --upgrade chex

Revision Not Found for url: https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/flax/vae/config.json. 是因为链接失效了，作者的回复如下链接: link，需要修改susie/susie/model.py文件

OSError: Can’t load the model in PyTorch format because PyTorch is not installed. Please, install PyTorch or use native Flax weights.没下载torch

pip install torch==2.1.0+cu121 torchvision==0.16.0+cu121 -f https://download.pytorch.org/whl/torch_stable.html

6. 运行结果

这是运行的示例结果，后续还会对这篇论文的训练过程以及在CALVIN环境中的实验和训练过程进行复现记录。