siliconflow / onediff

OneDiff: An out-of-the-box acceleration library for diffusion models.
https://github.com/siliconflow/onediff/wiki
Apache License 2.0
1.7k stars 104 forks source link

sd-webui refactor, and support refiner model #930

Closed marigoold closed 5 months ago

marigoold commented 5 months ago

重构了 webui 插件,主要是把相同的代码函数挪到一起;支持了 refiner,之前用 refiner 会报错,并且第二次推理 base 不会加速。

关于 refiner: refiner 其实可以看做是一个 unet,在推理的某个时刻替换 base 模型的 unet,使图像更高清。使用 base + refiner 推理的过程可以(不严谨地)看作是两个模型,在 n 个 step 的前 m 个 step 用第一个模型推理,剩下的 step 用第二个模型推理。

之前使用 onediff 跑 base + refiner 会报错,即使能跑起来也没有加速作用,原因如下:

  1. 第一次推理正常,第二次推理时候报错 meta tensor 不能转化成 oneflow tensor。因为 webui 在加载 refiner 后,会把 base 销毁(也比较合理,因为之后的推理 base 模型不会再使用了),销毁的做法是 model.to(meta),所以在第二次推理时,拿到的 unet 是个 meta model,就报错了
  2. 解决了1之后,发现第二次推理时不会加速,原因是,在第二次推理之前,unet 默认是 refiner,在进入 process_images 之后切换成 base,但是 oneflow 把 unet 设置成 graph 是在进入这个函数之前。所以时间线是这样的 unet是refiner -> oneflow 把 unet 替换成 graph -> webui 加载 base,再把 unet 替换成 base 这就导致用 base 推理时,base 是最后一步的 webui 加载进来的纯 torch model,没有加速效果 本 PR 解决了上面两个问题。

注意:

  1. 目前支持 refiner 的情况是,refiner 部分还是用 eager 跑,这么做有两个考虑的原因:
    • refiner 会在推理时临时加载,而且作用的 steps 通常不多。相比于加载模型的时间开销,编译节省的时间开销可能不太显著
    • 目前 base 模型需要驻留在内存中(不然 graph 会失效),如果 refiner 也编译,可能显存会吃不消
  2. base + refiner 一起跑峰值显存会比 torch 原生高很多(主要就是一个 base 模型的大小)