TylerYep / torchinfo

View model summaries in PyTorch!

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

The parameter of stable diffusion unet is wrong

FlyHighest opened this issue · comments

Describe the bug
I'm using torchinfo summary to see the number of parameters of UNet in Stable-Diffusion-2, but find it unmatch what is reported publicly.

To Reproduce

from diffusers import StableDiffusionPipeline
pipe=StableDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1')
pipe.to('cuda')
from torchinfo import summary
summary(pipeline.unet,[[1,4,32,32],[1],[1,77,1024]])

Expected behavior
The total parameter should be 0.8B. But it shows 2B.

Additional context
The torchinfo is 1.8.0. diffusers is 0.24.0.