JDiffusion: Jittor的扩散模型库发布了！

2024/02/08 Tutorial

JDiffusion: Jittor的扩散模型库发布了！

2月5日，计图团队发布了Jittor的扩散模型库，可支持多种扩散模型，例如StableDiffusion、SD-XL、ControlNet、LCM等前沿模型。

计图(Jittor)是清华大学计算机系图形学实验室于2020年3月20日发布并开源的深度学习框架[1]。计图团队已经陆续发布了GAN模型库、语义分割模型库、检测与实例分割库、三维点云库、可微渲染库、遥感检测库、NeRF模型库、稀疏卷积计算库、大模型推理库等。

Part 1

背景介绍

近年来，扩散模型在图片、音频、视频等等领域发展迅速，已经成为目前学术界乃至行业常用的AIGC生成模型的首要选择。

例如，Stable Diffusion系列模型已经成为学术界主流的基础模型，被非常多的图片生成工作使用，具有相当大的影响力。

图1 StableDiffusion-XL模型生成图

除此之外，Diffusers模型库作为目前最大的扩散模型开源库，每个月都有百万的上传与下载量，被数千项目使用，是目前主要的扩散模型库。
Part 2

JDiffusion库介绍

JDiffusion库是基于Jittor、依托JTorch、参考Diffusers框架开发的一个全面兼容huggingface的扩散模型库，可以很好地兼容与支持目前的各种开源扩散模型。

JDiffusion能够完美衔接到当前的社区生态，很容易兼容目前主流社区支持的开源模型，我们采取了非常易用的实现，方便研究人员编写代码与测试。

JDiffusion目前已经支持了StableDiffusion[2], SD-XL[3], instruct-pix2pix[4], ControlNet[5], AnimateDiff[6], LCM[7]等多种扩散模型的推理。以下是一些示例。

SDXL
from JDiffusion import StableDiffusionXLPipelineimport jittor as jtjt.flags.use_cuda = 1pipeline_text2image = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", use_safetensors=True)prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"image = pipeline_text2image(prompt=prompt).images[0]image.save("./output/test_sdxl_text2img.jpg")

通过这个简易的代码我们可以获得如下精美的图片：

图2 SD-XL根据上述prompt得到的图片

Instruct-pix2pix
from PIL import Imagefrom JDiffusion.pipelines.pipeline_stable_diffusion_instruct_pix2pix import StableDiffusionInstructPix2PixPipelineimport jittor as jtjt.flags.use_cuda = 1image = Image.open("./asset/mountain.png").convert("RGB").resize((512, 512))pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained("timbrooks/instruct-pix2pix", dtype=jt.float32)prompt = "make the mountains snowy"image = pipe(prompt=prompt, image=image).images[0]image.save("./output/test_ip2p.jpg")
通过这段代码，可生成如下的图片风格迁移。
图3 Instruct-pix2pix修改前后的图片

ControlNet

from PIL import Imageimport cv2import numpy as npimport jittor as jtjt.flags.use_cuda = 1image = np.array(Image.open("./asset/input_image_vermeer.png"))image = cv2.Canny(image, 100, 200)[:, :, None]image = np.concatenate([image, image, image], axis=2)canny_image = Image.fromarray(image)from JDiffusion import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepSchedulercontrolnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", use_safetensors=True)pipe = StableDiffusionControlNetPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", controlnet=controlnet, use_safetensors=True)pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)output = pipe("the mona lisa", image=canny_image).images[0]output.save("./output/test_controlnet.png")
通过这段代码，可生成如下的图片效果。

图4 ControlNet运行前后效果图

AnimateDiff
import jittor as jtjt.flags.use_cuda = 1from JDiffusion.pipelines import AnimateDiffPipelinefrom diffusers import DDIMScheduler, MotionAdapterfrom diffusers.utils import export_to_gifmodel_id = "SG161222/Realistic_Vision_V5.1_noVAE"adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2")pipe = AnimateDiffPipeline.from_pretrained(model_id, motion_adapter=adapter)scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler", clip_sample=False, timestep_spacing="linspace", beta_schedule="linear", steps_offset=1,)pipe.scheduler = schedulerframes = pipe(prompt=("masterpiece, bestquality, highlydetailed, ultradetailed, sunset, moon on the sky""orange sky, warm lighting, fishing boats, ocean waves seagulls, "        "rippling water, wharf, silhouette, serene atmosphere, dusk, evening glow, ""golden hour, coastal landscape, seaside scenery"),    negative_prompt="bad quality, worse quality", num_frames=16, guidance_scale=7.5, num_inference_steps=25).frames[0]export_to_gif(frames, "./output/animation.gif")
通过这段代码，可生成如下的动画效果。

图5 AnimateDiff的效果

更多运行结果和代码可以到我们的仓库示例上查看：

https://github.com/JittorRepos/JDiffusion/tree/master/examples

除此之外，目前JDiffusion代码以及使用方法已经全部开源到Github上。

https://github.com/JittorRepos/JDiffusion

欢迎大家使用JDiffusion开展研究工作！

后续Jittor框架将支持更多Diffusion模型，也将支持扩散模型的训练，欢迎大家参与进来，一起助力Jittor开源生态！

参考文献
  1. Shi-Min Hu, Dun Liang, Guo-Ye Yang, Guo-Wei Yang and Wen-Yang Zhou, Jittor: a novel deep learning framework with meta-operators and unified graph execution, Science China Information Science, 2020, Vol. 63, No.12, article no. 222103.
  2. Rombach R, Blattmann A, Lorenz D, et al., High-resolution image synthesis with latent diffusion models，CVPR, 10684-10695，2022.
  3. Podell D, English Z, Lacey K, et al., Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv: 2307.01952, 2023.
  4. Brooks T, Holynski A, Efros A A，Instructpix2pix: Learning to follow image editing instructions, CVPR, 18392-18402, 2023.
  5. Zhang L, Rao A, Agrawala M, Adding conditional control to text-to-image diffusion models, CVPR, 3836-3847, 2023.
  6. Guo Y, Yang C, Rao A, et al., Animatediff: Animate your personalized text-to-image diffusion models without specific tuning. arXiv preprint arXiv:2307.04725, 2023.
  7. Luo S, Tan Y, Huang L, et al., Latent consistency models: Synthesizing high-resolution images with few-step inference, arXiv preprint arXiv:2310.04378, 2023.

GGC往期回顾

1. 全面升级！计图更新用户文档

2. 计图助力非十科技发布AI代码助手，代码大模型速度精度超越Copilot

3. Computational Visual Media第10卷第1期导读

4. 五位图形学领域学者入选IEEE Fellow

5. CVMJ获最新期刊影响因子6.9，计算机图形学领域排名第一

可通过下方二维码，关注清华大学图形学实验室，了解图形学、Jittor框架、CVMJ期刊和CVM会议的相关资讯。

Search

    Table of Contents