Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

sd-webui refactor, and support refiner model #930

Merged
merged 19 commits into from
Jun 17, 2024

Conversation

marigoold
Copy link
Contributor

@marigoold marigoold commented Jun 4, 2024

重构了 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 模型的大小)

@marigoold marigoold marked this pull request as ready for review June 4, 2024 09:02
@lijunliangTG lijunliangTG self-requested a review June 17, 2024 07:56
@marigoold marigoold merged commit 156d4f0 into main Jun 17, 2024
7 checks passed
@marigoold marigoold deleted the webui-refactor-support-refiner branch June 17, 2024 07:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants