search

facebookincubator / AITemplate

AITemplate is a Python framework which renders neural network into high performance CUDA/HIP C++ code. Specialized for FP16 TensorCore (NVIDIA GPU) and MatrixCore (AMD GPU) inference.
Apache License 2.0
4.52k stars 360 forks source link

How do I make stable diffusion with AIT work with any latent resolution? #867

Open comfyanonymous opened 1 year ago

comfyanonymous commented 1 year ago

The original stable diffusion implementation for SD1.x only supports latent sizes that are multiples of 8 however there is a trick (here it is implemented in diffusers) to upsample the latent in the upsample block to the exact size so that the later concat doesn't fail. This makes it work with any latent resolution.

I would like to implement that here: https://github.com/facebookincubator/AITemplate/blob/main/examples/05_stable_diffusion/src/modeling/resnet.py#L64

If it was possible to slice the tensor like this after the upsampling it would work:

            x = nn.Upsampling2d(scale_factor=2.0, mode="nearest")(x)
            b, h, w, c = output_size
            x = x[:, :h, :w]

However ops.dynamic_slice does not work when the size parameters are ranges. Is there a way I can do this that will work with dynamic sizes?

Another way would be to make the nn.Upsampling2d op take the output size as input instead of the scale factor: The kernel seems to take an output size instead of a scale factor so maybe that would be easy to add.

What's the best way to solve this problem?

hlky commented 1 year ago

@comfyanonymous I can confirm nn.Upsampling2d can be made to work with explicit output size rather than scale factor. I will bring the changes here soon :)

Shaistrong commented 1 year ago

@comfyanonymous I can confirm nn.Upsampling2d can be made to work with explicit output size rather than scale factor. I will bring the changes here soon :)

Hey, I'm also working on this; how did you guys compile the Windows AIT modules?

hlky commented 1 year ago

@Shaistrong #841 is not merged yet, you can use a local branch with it applied. VS2022, x64 Native Tools Command Prompt for VS 2022, set AIT_USE_CMAKE_COMPILATION=1, then run compilation scripts. Comment on the PR if you have any issues related to MSVC, I may already know a fix.

Shaistrong commented 1 year ago

@Shaistrong #841 is not merged yet, you can use a local branch with it applied. VS2022, x64 Native Tools Command Prompt for VS 2022, set AIT_USE_CMAKE_COMPILATION=1, then run compilation scripts. Comment on the PR if you have any issues related to MSVC, I may already know a fix.

fatal error C1083: Cannot open include file: 'dlfcn.h': No such file or directory any idea what this means?

Shaistrong commented 1 year ago

@Shaistrong #841 is not merged yet, you can use a local branch with it applied. VS2022, x64 Native Tools Command Prompt for VS 2022, set AIT_USE_CMAKE_COMPILATION=1, then run compilation scripts. Comment on the PR if you have any issues related to MSVC, I may already know a fix.

also; spits at me a huge paragraph of errors of my directories, then it shows that error I quoted on my last update