• hidenorly's avatar
    Add FP32 fallback support on sd_vae_approx · 58c19545
    hidenorly authored
    This tries to execute interpolate with FP32 if it failed.
    
    Background is that
    on some environment such as Mx chip MacOS devices, we get error as follows:
    
    ```
    "torch/nn/functional.py", line 3931, in interpolate
            return torch._C._nn.upsample_nearest2d(input, output_size, scale_factors)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        RuntimeError: "upsample_nearest2d_channels_last" not implemented for 'Half'
    ```
    
    In this case, ```--no-half``` doesn't help to solve. Therefore this commits add the FP32 fallback execution to solve it.
    
    Note that the submodule may require additional modifications. The following is the example modification on the other submodule.
    
    ```repositories/stable-diffusion-stability-ai/ldm/modules/diffusionmodules/openaimodel.py
    
    class Upsample(nn.Module):
    ..snip..
        def forward(self, x):
            assert x.shape[1] == self.channels
            if self.dims == 3:
                x = F.interpolate(
                    x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
                )
            else:
                try:
                    x = F.interpolate(x, scale_factor=2, mode="nearest")
                except:
                    x = F.interpolate(x.to(th.float32), scale_factor=2, mode="nearest").to(x.dtype)
            if self.use_conv:
                x = self.conv(x)
            return x
    ..snip..
    ```
    
    You can see the FP32 fallback execution as same as sd_vae_approx.py.
    58c19545
sd_vae_approx.py 3.28 KB