# DragNoise
DragNoise 模型，利用扩散模型进行基于点的交互式图像编辑，强大而快速。

## 论文
`Drag Your Noise: Interactive Point-based Editing via Diffusion Semantic Propagation`
- https://arxiv.org/abs/2404.01050
- CVPR 2024

## 模型结构
<!-- 此处一句话简要介绍模型结构 -->
DragNoise利用扩散模型进行基于点的交互式图像编辑，算法整体思路沿着 DragDiffusion 算法：

<div align=center>
    <img src="./doc/overview.png"/>
    <div >DragNoise</div>
</div>


## 算法原理
<!-- 我们介绍了Dragnoise，提供了强大而加速的编辑，而无需重试潜在地图。Dragnoise的核心原理在于利用每个U-NET作为语义编辑器的预测噪声输出。这种方法以两个关键观察为基础：首先，U-NET的瓶颈本质上具有具有互动编辑的理想的语义丰富特征；其次，高级语义是在脱诺过程早期建立的，在后续阶段显示出最小的差异。利用这些见解，在单个降解步骤中进行了Dragnoise编辑扩散语义，并有效地传播了这些变化，从而确保了扩散编辑的稳定性和效率。 -->

DragNoise 算法整体思路沿着 DragDiffusion 算法，与此不同的是针对扩散模型中的 “middle-block replacement” 的操作进行探索。该操作从某个去噪时间步开始，将不同层的特征复制到所有后续 timestep 的对应层。通过观察 DDIM inversion 重建图像的效果，探索扩散模型在何时以及何处学习到何种层次的语义信息。 

通过实验发现，bottleneck 特征是一种最优扩散语义表示，适合于高效编辑。由于它可以在早期 timestep 中有效地被编辑，因此操纵 bottleneck 特征可以平滑地传播到后面的去噪步骤，从而确保结果图像扩散语义的完整性。此外，由于优化 bottleneck 的路径短，有效地避免了梯度消失问题。

<!-- <div align=center>
    <img src="./doc/pipeline.png"/>
    <div >DragNoise</div>
</div> -->


## 环境配置
```
mv dragnoise_pytorch dragnoise # 去框架名后缀
# docker 的 -v 路径、docker_name 和 imageID 根据实际情况修改
# pip 安装时如果出现下载慢可以尝试别的镜像源
```
### Docker（方法一）
<!-- 此处提供[光源](https://www.sourcefind.cn/#/service-details)拉取docker镜像的地址与使用步骤 -->
```
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.2-py3.10 # 本镜像imageID为：2f1f619d0182
docker run -it -v /path/your_code_data/:/path/your_code_data/ -v /opt/hyhal/:/opt/hyhal/:ro --shm-size=16G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --network=host --name docker_name imageID bash
cd /your_code_path/dragnoise
pip install -r requirements.txt
```
### Dockerfile（方法二）
<!-- 此处提供dockerfile的使用方法 -->
```
cd /your_code_path/dragnoise/docker
docker build --no-cache -t codestral:latest .
docker run -it -v /path/your_code_data/:/path/your_code_data/ -v /opt/hyhal/:/opt/hyhal/:ro --shm-size=16G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --network=host --name docker_name imageID bash
cd /your_code_path/dragnoise
pip install -r requirements.txt
```
### Anaconda（方法三）
<!-- 此处提供本地配置、编译的详细步骤，例如： -->

关于本项目DCU显卡所需的特殊深度学习库可从[光合](https://developer.sourcefind.cn/tool/)开发者社区下载安装。
```
DTK驱动: dtk24.04.2
python: python3.10
pytorch: 2.1.0
```
`Tips：以上DTK驱动、python、pytorch等DCU相关工具版本需要严格一一对应`

其它非深度学习库参照requirements.txt安装：
```
pip install -r requirements.txt
```
## 数据集
测试数据集 [DragBench](https://github.com/Yujun-Shi/DragDiffusion/releases/download/v0.1.1/DragBench.zip)。\
下载后放在 ./drag_bench_evaluation/drag_bench_data 并解压，文件构成:
<br>
DragBench<br>
--- animals<br>
------ JH_2023-09-14-1820-16<br>
------ JH_2023-09-14-1821-23<br>
------ JH_2023-09-14-1821-58<br>
------ ...<br>
--- art_work<br>
--- building_city_view<br>
--- ...<br>
--- other_objects<br>
<br>



## 训练
```
# LoRA微调
python run_lora_training.py
# 训练drag扩散模型，并输出结果
python run_drag_diffusion.py
# 如果出现huggingface访问不通，请执行 `export HF_ENDPOINT=https://hf-mirror.com`
```
亦或者webui界面。

## 推理
<!-- 下载模型权重:
```
python scripts/download_model.py
```
或者从 [SCNet](http://113.200.138.88:18080/aimodels/findsource-dependency/stylegan2_pytorch) 上快速下载，并放在 /checkpoints 文件夹下。 -->

可视化webui推理：
```
python drag_ui.py --listen
# 如果出现huggingface访问不通，请执行 `export HF_ENDPOINT=https://hf-mirror.com`
```
<div align=center>
    <img src="./doc/webui.jpg" width=600/>
    <div >webui界面</div>
</div>

1、上传图片；\
2、输入提示；\
3、LoRA训练；\
4、通过鼠标选择要编辑的区域；\
5、通过鼠标标记点位；\
6、运行。\
ps：Drag以及LoRA的一些参数自行视情况修改。

## result
<!-- 此处填算法效果测试图（包括输入、输出） -->

<div align=center>
    <img src="./doc/input.jpg" width=600/>
    <div >输入</div>
</div>


<div align=center>
    <img src="./doc/result.png" width=600/>
    <div >推理结果</div>
</div>

<!-- <div align=center>
    <img src="./doc/image (1).png" width=600/>
    <div >输出</div>
</div> -->


### 精度
测试集 `DragBench`，如上所述下载并解压好。
```
python run_lora_training.py
python run_drag_diffusion.py
python run_eval_similarity.py
python run_eval_point_matching.py
# ps:上述脚本的一些文件路径自行根据情况修改
```

| 加速卡K100_AI | 1-lpips ↑ | clip sim ↑| mean distance ↓ |
| :-----| :----- | :---- | :---- |
| paper | 0.894 | 0.971 | 33.404 |
| 优化后 | 0.885 | 0.967 | 30.206 |
<!-- | 单元格 | 单元格 | 单元格 | -->

ps:优化两点 1）lora训练部分的rank选择；2）drag部分采用了多层融合。

<!-- <div align=center>
    <img src="./doc/result2.jpg" width=600/>
</div>
<div align=center>
    <img src="./doc/result3.jpg" width=600/>
</div>
<div align=center>
    <img src="./doc/result4.jpg" width=600/>
</div>
<div align=center>
    <img src="./doc/result5.jpg" width=600/>
    <div >paper结果和优化后的结果</div>
</div> -->

## 应用场景
### 算法类别

<!-- 超出以上分类的类别命名也可参考此网址中的类别名：https://huggingface.co/ \ -->
`AIGC`

### 热点应用行业
<!-- 应用行业的填写需要做大量调研，从而为使用者提供专业、全面的推荐，除特殊算法，通常推荐数量>=3。 -->
`零售,制造,电商,医疗,教育`

<!-- ## 预训练权重 -->
<!-- - 此处填写预训练权重在公司内部的下载地址（预训练权重存放中心为：[SCNet AIModels](http://113.200.138.88:18080/aimodels) ，模型用到的各预训练权重请分别填上具体地址。），过小权重文件可打包到项目里。
- 此处填写公开预训练权重官网下载地址（非必须）。 -->

## 源码仓库及问题反馈
<!-- - 此处填本项目gitlab地址 -->
- https://developer.sourcefind.cn/codes/modelzoo/dragnoise_pytorch
## 参考资料
- https://github.com/haofengl/DragNoise