img2img-turbo简介
img2img-turbo 的框架结构设计紧凑、灵活,主要围绕着扩散模型和生成对抗网络(GAN)的优化实现。其架构设计有助于在图像到图像转换任务中实现高效的推理和良好的生成效果,尤其是在不同场景之间保持输入图像结构的完整性。以下是该框架的主要组成部分:
1. 核心模型结构
- 扩散模型与轻量化调优:img2img-turbo 通过改进预训练的扩散模型(如 Stable Diffusion),并引入 LoRA(Low-Rank Adaptation)适配器,使模型更加高效。LoRA 的使用帮助减少模型的参数量,同时通过保持某些层冻结,减少了对计算资源的需求。
- 跳跃连接与 Zero-Convs:为了更好地在图像生成过程中保持图像结构,模型使用了跳跃连接和零卷积层(Zero-Convs),这确保了输入图像和输出图像之间的重要信息传递。
2. 主要模型模块
- pix2pix-turbo:这是项目中的成对图像翻译模块,专门用于将输入图像转换为与其配对的输出图像。例如,它可以将草图转换为真实图片,或将边缘图像转换为完整的场景。
- CycleGAN-Turbo:这是一个非成对图像翻译模型,能够在没有成对训练数据的情况下进行图像转换任务,如白天转换为夜晚,或者天气从晴天变为雨天。它在无监督学习场景中表现尤其突出。
3. 训练与推理过程
- 快速推理:img2img-turbo 通过单步推理来加速图像生成过程,相比传统的逐步迭代扩散模型,这种优化显著减少了推理时间。例如,512x512 尺寸图像在 A100 GPU 上的推理时间仅为 0.11 秒。
- 多种任务支持:该框架支持多种图像转换任务,如 Edge2Image(边缘到图像)、Sketch2Image(草图到图像) 和 Day2Night(日夜转换),这些任务可以通过简单的调整模型架构或输入数据来实现。
4. 环境设置与灵活扩展
- 环境与依赖:该项目提供了 conda 和 virtual environment 两种环境配置方式,方便用户快速搭建和运行模型。其训练和推理代码可以直接在本地或远程服务器上运行。
- Gradio 集成:为了便于用户体验,img2img-turbo 提供了集成的 Gradio 界面,用户可以通过图形界面直接进行图像转换操作,而不需要深入编写代码。
5. 性能优化
- 跳跃连接与 LoRA 模块的集成 让模型在执行任务时既能保持高效推理,又能在转换过程中保留输入图像的关键特征。与其他模型相比,这种优化使其在图像到图像转换任务中能实现更高的保真度和更快的响应时间。
通过这些设计,img2img-turbo 实现了在多任务处理、快速推理、环境适配性和资源效率之间的良好平衡,适合在创意设计、科研、环境监测等多个领域中使用。
部署img2img-turbo
相关链接
https://github.com/GaParmar/img2img-turbo
https://zhuanlan.zhihu.com/p/688973522
https://blog.youkuaiyun.com/gitblog_00320/article/details/141050129
img2img-turbo: 革命性的单步图像转换技术 - 懂AI
在线sketch2image测试平台(国内无法访问):
Img2img Turbo Sketch - a Hugging Face Space by gparmar
部署环境
conda env:
conda env create -f environment.yaml
conda activate img2img-turbo
注意每个包的版本问题
遇到的问题
conda env create
部署环境时由于网速的问题,创建不了conda env环境。解决方法就是多试几次,总会有一次成功的。(网络问题)
huggingface_hub
版本:0.26
ImportError: cannot import name ‘cached_download’ from ‘huggingface_hub’ (/home/hxr/anaconda3/envs/img2img-turbo/lib/python3.10/site-packages/huggingface_hub/init.py)
升级,但是又出现以下diffusers的问题
diffusers
版本:0.30.3
AttributeError: ‘AutoencoderKL’ object has no attribute ‘add_adapter’
解决方法:huggingface_hub 安装为0.25以下,diffusers版本为0.25.1,可以跑代码。
模型加载失败
RuntimeError: PytorchStreamReader failed reading zip archive: failed finding central directory
解决方法:删除checkpoint中已下载的模型,由于某种原因模型下载不完整或者损坏,所以重新下载模型。模型下载地址:https://www.cs.cmu.edu/~img2img-turbo/models/
测试命令
paired
python src/inference_paired.py --model_name “edge_to_image” --input_image “assets/examples/bird.png” --prompt “a blue bird” --output_dir “outputs”
python src/inference_paired.py --model_name “sketch_to_image_stochastic” --input_image “assets/examples/sketch_input.png” --gamma 0.4 --prompt “ethereal fantasy concept art of an asteroid. magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy” --output_dir “outputs”
unpaired
python src/inference_unpaired.py --model_name “day_to_night” --input_image “assets/examples/day2night_input.png” --output_dir “outputs”
python src/inference_unpaired.py --model_name “night_to_day” --input_image “assets/examples/night2day_input.png” --output_dir “outputs”
python src/inference_unpaired.py --model_name “clear_to_rainy” --input_image “assets/examples/clear2rainy_input.png” --output_dir “outputs”
python src/inference_unpaired.py --model_name “rainy_to_clear” --input_image “assets/examples/rainy2clear_input.png” --output_dir “outputs”
gradio可视化界面
gradio gradio_sketch2image.py
gradio gradio_canny2image.py
端口从服务器映射到自己的物理主机:
ssh -L 7860:127.0.0.1:7860 username@ip
训练自己的模型img2mig-turbo
训练paired data(pix2pix-turbo)
第一步 获取数据集
bash scripts/download_fill50k.sh
第二步 训练模型
accelerate config
配置文件保存在/home/hxr/.cache/huggingface/accelerate/default_config.yaml
第三步 监控训练过程
使用wandb来监控训练过程:Weights & Biases: The AI Developer Platform
- pip install wandb
- wandb login
- 输出wandb官网的API
wandb: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)
wandb: You can find your API key in your browser here: https://wandb.ai/authorize
wandb: Paste an API key from your profile and hit enter, or press ctrl+c to quit:
wandb: Appending key for api.wandb.ai to your netrc file: /home/hxr/.netrc
数据集:my_fill50k
训练模型命令:accelerate launch src/train_pix2pix_turbo.py --pretrained_model_name_or_path=“stabilityai/sd-turbo” --output_dir=“output/pix2pix_turbo/fill50k” --dataset_folder=“data/my_fill50k” --resolution=256 --train_batch_size=1 --enable_xformers_memory_efficient_attention --viz_freq 25 --track_val_fid --report_to “wandb” --tracker_project_name “pix2pix_turbo_fill50k”
accelerate launch src/train_pix2pix_turbo.py --pretrained_model_name_or_path=“stabilityai/sd-turbo” --output_dir=“output/pix2pix_turbo/fill50k” --dataset_folder=“data/my_fill50k” --resolution=512 --train_batch_size=2 --viz_freq 25 --track_val_fid --report_to “wandb” --tracker_project_name “pix2pix_turbo_fill50k”
accelerate launch src/train_pix2pix_turbo.py --pretrained_model_name_or_path=“stabilityai/sd-turbo” --output_dir=“output/pix2pix_turbo/fill50k” --dataset_folder=“data/my_fill50k” --resolution=256 --train_batch_size=1 --enable_xformers_memory_efficient_attention --viz_freq 20 --track_val_fid --report_to “wandb” --tracker_project_name “pix2pix_turbo_fill50k”
数据集:facades
accelerate launch src/train_pix2pix_turbo.py --pretrained_model_name_or_path=“stabilityai/sd-turbo” --output_dir=“output/pix2pix_turbo/facades” --dataset_folder=“data/facades” --resolution=512 --train_batch_size=2 --enable_xformers_memory_efficient_attention --viz_freq 25 --track_val_fid --report_to “wandb” --tracker_project_name “pix2pix_turbo_facades”
遇到的问题
1.RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
这个错误 RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0
表示你的代码在执行某个操作时,使用的张量(tensor)没有在同一个设备上,有的在 CPU 上,有的在 GPU(cuda:0)上,而 PyTorch 需要所有操作都在同一个设备上进行。
**解决办法:**重新accelerate config,将其设置正确。参考链接:https://blog.youkuaiyun.com/zjc910997316/article/details/134642194
2.torch.cuda.OutOfMemoryError: CUDA out of memory.
原因:实验室服务器只有一块RTX 4090服务器,算力和gpu内存实在有限。
GPU显存放不下训练的模型。
3.wandb 连接失败
wandb: ERROR Run initialization has timed out after 90.0 sec.
这个问题是由于这个训练可视化平台的网络连接问题
解决方法:
第四步 使用训练好的模型运行推理
python src/inference_paired.py --model_path “output/pix2pix_turbo/fill50k/checkpoints/model_6001.pkl” --input_image “data/my_fill50k/test_A/40000.png” --prompt “violet circle with orange background” --output_dir “outputs”
训练 Unpaired Data(CycleGAN-turbo)
第一步:bash scripts/download_horse2zebra.sh
第二步:accelerate config
第三步:
export NCCL_P2P_DISABLE=1
accelerate launch --main_process_port 29501 src/train_cyclegan_turbo.py --pretrained_model_name_or_path=“stabilityai/sd-turbo” --output_dir=“output/cyclegan_turbo/my_horse2zebra” --dataset_folder “data/my_horse2zebra” --train_img_prep “resize_286_randomcrop_256x256_hflip” --val_img_prep “no_resize” --learning_rate=“1e-5” --max_train_steps=25000 --train_batch_size=1 --gradient_accumulation_steps=1 --report_to “wandb” --tracker_project_name “gparmar_unpaired_h2z_cycle_debug_v2” --enable_xformers_memory_efficient_attention --validation_steps 250 --lambda_gan 0.5 --lambda_idt 1 --lambda_cycle 1
使用训练好的模型进行图像转换任务
python src/inference_unpaired.py --model_path “output/cyclegan_turbo/my_horse2zebra/checkpoints/model_2501.pkl” --input_image “data/my_horse2zebra/test_A/n02381460_20.jpg” --prompt “picture of a zebra” --direction “a2b” --output_dir “outputs” --image_prep “no_resize”
对比实验
pix2pix
CycleGAN
pytorch-pix2pix and CycleGAN
https://github.com/junyanz/CycleGAN
https://github.com/phillipi/pix2pix
https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
https://github.com/lllyasviel/ControlNet
https://dl.cv.ethz.ch/bdd100k/data/
https://github.com/Janspiry/Palette-Image-to-Image-Diffusion-Models
img2img-turbo工作流程
-
模型初始化:
- 使用
model.py
加载并配置 SD-Turbo 模型,包括 Encoder、UNet 和 Decoder。配置 LoRA 权重和跳跃连接,为后续任务适配和细节保留奠定基础。
- 使用
-
数据准备:
- 使用
image_prep.py
进行数据预处理,将输入数据(如边缘图或草图)调整为模型所需的格式,为模型的配对和无配对任务准备输入图像。
- 使用
-
训练阶段:
-
配对任务训练
(如边缘到图像、草图到图像):
- 使用
pix2pix_turbo.py
和train_pix2pix_turbo.py
进行训练,通过重构损失和 GAN 损失等来提高图像生成的质量。
- 使用
-
无配对任务训练
(如昼夜转换、天气变化):
- 使用
cyclegan_turbo.py
和train_cyclegan_turbo.py
,引入循环一致性损失和对抗性损失来训练模型,无需配对数据。
- 使用
-
-
推理与生成:
- 使用
inference_paired.py
处理配对任务推理,实现高质量的图像转换(如 Sketch2Image)。 - 使用
inference_unpaired.py
进行无配对任务推理,实现昼夜或天气变化等场景的转换。
- 使用
-
用户交互界面:
- 使用
gradio_canny2image.py
和gradio_sketch2image.py
构建用户界面,通过 Gradio 提供边缘到图像、草图到图像等实时转换功能,便于用户直接体验模型效果。
- 使用
这一流程逐步展示了 img2img-turbo 项目的工作顺序,从模型配置到用户交互,实现完整的图像到图像转换任务。
Kaggle 部署
!git clone https://github.com/GaParmar/img2img-turbo
! pip install -r /kaggle/working/img2img-turbo/requirements.txt
! python /kaggle/working/img2img-turbo/src/inference_paired.py --model_name “edge_to_image” --input_image “/kaggle/working/img2img-turbo/assets/examples/bird.png” --prompt “a blue bird” --output_dir “outputs”
! python /kaggle/working/img2img-turbo/src/inference_unpaired.py --model_name “night_to_day” --input_image “/kaggle/working/img2img-turbo/assets/examples/night2day_input.png” --output_dir “outputs”