HiDiffusion

<div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/8ab4aafb-e193-4cad-8391-612466779ff0.jpg" height=120> </div>

<div align="center">💡 HiDiffusion: 释放预训练扩散模型更高分辨率的创造力和效率</div>

<div align="center">Shen Zhang, Zhaowei Chen, Zhenyu Zhao, Yuhao Chen, Yao Tang, Jiajun Liang</div> <br> <div align="center"> <a href="https://hidiffusion.github.io/"><img src="https://img.shields.io/static/v1?label=项目主页&message=Github&color=blue&logo=github-pages"></a> &ensp; <a href="https://arxiv.org/abs/2311.17528"><img src="https://img.shields.io/static/v1?label=论文&message=Arxiv:HiDiffusion&color=red&logo=arxiv"></a> &ensp; <a href="https://colab.research.google.com/drive/1EiBn9lSnPZTU4cikRRaBBexs429M-qty?usp=sharing"><img src="https://img.shields.io/static/v1?label=演示&message=Colab&color=purple&logo=googlecolab"></a> &ensp; <a href="https://openbayes.com/console/public/tutorials/SaPYcYCaWSA"><img src="https://img.shields.io/static/v1?label=演示&message=OpenBayes&color=green"></a> &ensp; </div> <div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/f8fe1892-7663-45c7-bd21-599e936a7a7a.jpg" width="800" ></img> <br> <em> （不同扩散模型、分辨率和宽高比的HiDiffusion样本选择。） </em> </div> <br>

👉 为什么选择HiDiffusion

一种无需训练即可提高预训练扩散模型分辨率和速度的方法。
设计为即插即用的实现。只需添加一行代码即可集成到扩散管道中！
支持各种任务，包括文本到图像、图像到图像、图像修复。

<div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/9aaf030f-e65d-49d1-ba0e-5d33bfd33db8.jpg" width="800" ></img> <br> <em> （更快速，更好的图像细节。） </em> </div> <br> <div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/46fc2b90-4fa5-4a70-b5e3-aa06400d0518.jpg" width="800" ></img> <br> <em> （ControlNet和图像修复任务的2K结果。） </em> </div> <br>

🔥 更新

2024.8.15 - 💥 Diffusers文档已添加HiDiffusion，详见此处。感谢Diffusers团队！
2024.7.3 - 💥 被ECCV 2024接收！
2024.6.19 - 💥 已集成到OpenBayes，查看演示。感谢OpenBayes团队！
2024.6.16 - 💥 支持PyTorch 2.X。
2024.6.16 - 💥 修复非方形图像生成问题。现在HiDiffusion支持更多图像尺寸和宽高比。
2024.5.7 - 💥 支持图像到图像任务，详见此处。
2024.4.16 - 💥 发布源代码。

📢 支持的模型

注意：HiDiffusion还支持基于这些仓库的下游扩散模型，如Ghibli-Diffusion、Playground等。

💣 支持的任务

✅ 文本到图像
✅ ControlNet，包括文本到图像、图像到图像
✅ 图像修复

🔎 主要依赖

本仓库在以下环境中测试通过

Python==3.8
torch>=1.13.1
diffusers>=0.25.0
transformers
accelerate
xformers

🔑 安装HiDiffusion

安装主要依赖中的包后，安装HiDiffusion：

pip3 install hidiffusion

从源码安装

或者，您可以从github源码安装。克隆仓库并安装：

git clone https://github.com/megvii-model/HiDiffusion.git
cd HiDiffusion
python3 setup.py install

🚀 使用方法

使用HiDiffusion生成输出非常简单，基于🤗 diffusers。您只需添加一行代码。

文本到图像生成

Stable Diffusion XL

from hidiffusion import apply_hidiffusion, remove_hidiffusion
from diffusers import StableDiffusionXLPipeline, DDIMScheduler
import torch
pretrain_model = "stabilityai/stable-diffusion-xl-base-1.0"
scheduler = DDIMScheduler.from_pretrained(pretrain_model, subfolder="scheduler")
pipe = StableDiffusionXLPipeline.from_pretrained(pretrain_model, scheduler = scheduler, torch_dtype=torch.float16, variant="fp16").to("cuda")

# # 可选。enable_xformers_memory_efficient_attention可以节省内存使用并提高推理速度。enable_model_cpu_offload和enable_vae_tiling可以节省内存使用。
# pipe.enable_xformers_memory_efficient_attention()
# pipe.enable_model_cpu_offload()
# pipe.enable_vae_tiling()

# 用一行代码应用hidiffusion。
apply_hidiffusion(pipe)

prompt = "站立在废墟中，一个石头傀儡苏醒了，藤蔓和花朵从它身体的缝隙中生长出来。"
negative_prompt = "模糊，丑陋，重复，画得很差的脸，变形，马赛克，人工痕迹，畸形的肢体"
image = pipe(prompt, guidance_scale=7.5, height=2048, width=2048, eta=1.0, negative_prompt=negative_prompt).images[0]
image.save(f"golem.jpg")

设置height = 4096，width = 4096，您可以得到4096x4096分辨率的输出。

Stable Diffusion XL Turbo

from hidiffusion import apply_hidiffusion, remove_hidiffusion
from diffusers import AutoPipelineForText2Image
import torch
pretrain_model = "stabilityai/sdxl-turbo"
pipe = AutoPipelineForText2Image.from_pretrained(pretrain_model, torch_dtype=torch.float16, variant="fp16").to('cuda')

# # 可选。enable_xformers_memory_efficient_attention可以节省内存使用并提高推理速度。enable_model_cpu_offload和enable_vae_tiling可以节省内存使用。
# pipe.enable_xformers_memory_efficient_attention()
# pipe.enable_model_cpu_offload()
# pipe.enable_vae_tiling()

# 用一行代码应用hidiffusion。
apply_hidiffusion(pipe)

prompt = "在神秘森林的深处，一只机器猫头鹰用夜视镜头般的眼睛监视着夜间的生物。"
image = pipe(prompt, num_inference_steps=4, height=1024, width=1024, guidance_scale=0.0).images[0]
image.save(f"./owl.jpg")

Stable Diffusion v2-1

from hidiffusion import apply_hidiffusion, remove_hidiffusion
from diffusers import DiffusionPipeline, DDIMScheduler
import torch
pretrain_model = "stabilityai/stable-diffusion-2-1-base"
scheduler = DDIMScheduler.from_pretrained(pretrain_model, subfolder="scheduler")
pipe = DiffusionPipeline.from_pretrained(pretrain_model, scheduler = scheduler, torch_dtype=torch.float16).to("cuda")

# # 可选。enable_xformers_memory_efficient_attention 可以节省内存并提高推理速度。enable_model_cpu_offload 和 enable_vae_tiling 可以节省内存使用。
# pipe.enable_xformers_memory_efficient_attention()
# pipe.enable_model_cpu_offload()
# pipe.enable_vae_tiling()

# 使用单行代码应用 hidiffusion。
apply_hidiffusion(pipe)

prompt = "一只可爱快乐的棕色边境牧羊犬坐在床上，高细节。"
negative_prompt = "丑陋，平铺，画面外，面部画得不好，多余的肢体，畸形，变形，身体超出画面，模糊，解剖结构不佳，模糊不清，人工痕迹，比例不当。"
image = pipe(prompt, guidance_scale=7.5, height=1024, width=1024, eta=1.0, negative_prompt=negative_prompt).images[0]
image.save(f"collie.jpg")

设置 height = 2048，width = 2048，您可以获得 2048x2048 分辨率的输出。

Stable Diffusion v1-5

from hidiffusion import apply_hidiffusion, remove_hidiffusion
from diffusers import DiffusionPipeline, DDIMScheduler
import torch
pretrain_model = "runwayml/stable-diffusion-v1-5"
scheduler = DDIMScheduler.from_pretrained(pretrain_model, subfolder="scheduler")
pipe = DiffusionPipeline.from_pretrained(pretrain_model, scheduler = scheduler, torch_dtype=torch.float16).to("cuda")

# # 可选。enable_xformers_memory_efficient_attention 可以节省内存并提高推理速度。enable_model_cpu_offload 和 enable_vae_tiling 可以节省内存使用。
# pipe.enable_xformers_memory_efficient_attention()
# pipe.enable_model_cpu_offload()
# pipe.enable_vae_tiling()

# 使用单行代码应用 hidiffusion。
apply_hidiffusion(pipe)

prompt = "厚重笔触，明亮色彩，一只奇异的狐狸，可爱，Q版卡哇伊。细致的毛发，超高细节，大而有反光的眼睛，童话风格，ArtStation，居中构图，完美构图，居中，鲜艳色彩，柔和色调，高细节，8K。"
negative_prompt = "丑陋，平铺，面部画得不好，画面外，畸形，变形，模糊，解剖结构不佳，模糊不清。"
image = pipe(prompt, guidance_scale=7.5, height=1024, width=1024, eta=1.0, negative_prompt=negative_prompt).images[0]
image.save(f"fox.jpg")

设置 height = 2048，width = 2048，您可以获得 2048x2048 分辨率的输出。

移除 HiDiffusion

如果您想移除 HiDiffusion，只需使用 remove_hidiffusion(pipe)。

ControlNet

文本到图像生成

from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, DDIMScheduler
import numpy as np
import torch
import cv2
from PIL import Image
from hidiffusion import apply_hidiffusion, remove_hidiffusion

# 加载 assets 文件中的 Yoshua_Bengio.jpg。
path = './assets/Yoshua_Bengio.jpg'
image = Image.open(path)
# 获取 Canny 边缘图像
image = np.array(image)
image = cv2.Canny(image, 100, 200)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
canny_image = Image.fromarray(image)

# 初始化模型和管道
controlnet_conditioning_scale = 0.5  # 推荐用于良好的泛化
controlnet = ControlNetModel.from_pretrained(
    "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16, variant="fp16"
)
scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", controlnet=controlnet, torch_dtype=torch.float16,
    scheduler = scheduler
)

# 使用单行代码应用 hidiffusion。
apply_hidiffusion(pipe)

pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()

prompt = "小丑，高面部细节，高细节，柔和色彩，8K"
negative_prompt = "模糊，丑陋，重复，画工差，变形，马赛克。"

image = pipe(
    prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, image=canny_image,
    height=2048, width=2048, guidance_scale=7.5, negative_prompt = negative_prompt, eta=1.0
).images[0]

image.save('joker.jpg')

图像到图像生成

import torch
import numpy as np
from PIL import Image
from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline, DDIMScheduler
from hidiffusion import apply_hidiffusion, remove_hidiffusion
import cv2 

controlnet = ControlNetModel.from_pretrained(
    "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")

pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=controlnet,
    scheduler = scheduler,
    torch_dtype=torch.float16,
).to("cuda")

# 使用单行代码应用 hidiffusion。
apply_hidiffusion(pipe)

pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()

path = './assets/lara.jpeg'
ori_image = Image.open(path)
# 获取 Canny 边缘图像
image = np.array(ori_image)
image = cv2.Canny(image, 50, 120)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
canny_image = Image.fromarray(image)

controlnet_conditioning_scale = 0.5  # 推荐用于良好的泛化
prompt = "棕发的劳拉·克罗夫特，穿着背心和棕色背包。房间昏暗，装饰古朴，地板有图案，墙壁右侧有拱形设计和深色区域，柔和色调，高细节，8K高清获奖作品"
negative_prompt = "曝光不足，手部画得不好，重复的手，曝光过度，糟糕的艺术，初学者，业余，抽象，畸形，变形，特写，奇怪的颜色，水印"

image = pipe(prompt,
    image=ori_image,
    control_image=canny_image,
    height=1536,
    width=2048,
    strength=0.99,
    num_inference_steps=50,
    controlnet_conditioning_scale=controlnet_conditioning_scale,
    guidance_scale=12.5,
    negative_prompt = negative_prompt,
    eta=1.0
).images[0]

image.save("lara.jpg")

图像修复

import torch
from diffusers import AutoPipelineForInpainting, DDIMScheduler
from diffusers.utils import load_image
from hidiffusion import apply_hidiffusion, remove_hidiffusion
from PIL import Image

scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
pipeline = AutoPipelineForInpainting.from_pretrained(
    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16", 
    scheduler=scheduler
)

# 使用单行代码应用hidiffusion。
apply_hidiffusion(pipeline)

pipeline.enable_model_cpu_offload()
# 如果未安装xFormers，请删除以下行
pipeline.enable_xformers_memory_efficient_attention()

# 加载基础图像和蒙版图像
img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-text2img.png"
init_image = load_image(img_url)
# 加载assets文件中的mask_image.jpg。
mask_image = Image.open("./assets/mask_image.png")

prompt = "一位蒸汽朋克探险家，戴着皮质飞行员帽和护目镜，手持黄铜望远镜，站在巨大的古树中间，树根与复杂的齿轮和管道交织在一起。"

negative_prompt = "模糊、丑陋、重复、画工差、变形、马赛克"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, height=2048, width=2048, strength=0.85, guidance_scale=12.5, negative_prompt = negative_prompt, eta=1.0).images[0]
image.save('steampunk_explorer.jpg')

集成到下游模型

HiDiffusion支持基于支持的模型的模型，如Ghibli-Diffusion、Playground等。

Ghibli-Diffusion

from diffusers import StableDiffusionPipeline
import torch
from hidiffusion import apply_hidiffusion, remove_hidiffusion

model_id = "nitrosocke/Ghibli-Diffusion"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# 使用单行代码应用hidiffusion。
apply_hidiffusion(pipe)

prompt = "吉卜力风格的金发魔法公主"
negative_prompt="模糊、丑陋、重复、脸部画工差、变形、马赛克、瑕疵、四肢不良"
image = pipe(prompt, height=1024, width=1024, eta=1.0, negative_prompt=negative_prompt).images[0]

image.save("./magical_princess.jpg")

Playground

from diffusers import DiffusionPipeline
import torch
from hidiffusion import apply_hidiffusion, remove_hidiffusion

pipe = DiffusionPipeline.from_pretrained(
    "playgroundai/playground-v2-1024px-aesthetic",
    torch_dtype=torch.float16,
    use_safetensors=True,
    add_watermarker=False,
    variant="fp16"
)
pipe.to("cuda")
pipe.enable_xformers_memory_efficient_attention()

# 使用单行代码应用hidiffusion。
apply_hidiffusion(pipe)

prompt = "小女孩骑自行车，在宫崎骏风格的美丽动画场景中：一个下雪的东京城市，蓝天中漂浮着巨大的宫崎骏式云朵，阳光明媚的城市雪景令人着迷，宫崎骏式的风景意境，日本艺术"
negative_prompt="模糊、丑陋、重复、画工差、变形、马赛克"
image  = pipe(prompt=prompt, guidance_scale=3.0, height=2048, width=2048, negative_prompt=negative_prompt).images[0]
image.save('girl.jpg')

注意：您可以将guidance scale从3.0调整到5.0，并设计适当的negative prompt以生成令人满意的结果。

🙏 致谢

此代码库基于tomesd和diffusers。感谢！

🎓 引用

@article{zhang2023hidiffusion,
  title={HiDiffusion: Unlocking Higher-Resolution Creativity and Efficiency in Pretrained Diffusion Models},
  author={Zhang, Shen and Chen, Zhaowei and Zhao, Zhenyu and Chen, Yuhao and Tang, Yao and Liang, Jiajun},
  journal={arXiv preprint arXiv:2311.17528},
  year={2023}
}