Fix full previews and--no-half-vae to work correctly with --upcast-sampling

[brkirch]

Describe what this pull request is trying to achieve.

Fixes full previews and --no-half-vae to work correctly with --upcast-sampling. Also corrects a typo in CondFunc.

Additional notes and description of your changes

Casts the input to decode_first_stage and encode_first_stage to the VAE dtype. For encode_first_stage the input was previously incorrectly casted to the UNet dtype which can result in dtype mismatch errors with --no-half-vae. For decode_first_stage the input was previously not casted to the correct dtype for full previews, resulting in dtype mismatch errors without --no-half-vae.

Environment this was tested in

• OS: macOS
• Browser: Safari
• Graphics card: M1 Max 64 GB

[AUTOMATIC1111]

x = model.decode_first_stage(x.to(devices.dtype_vae) if devices.unet_needs_upcast else x)

Is this right?

Also I wanted to test it but never got to it, are all those if devices.unet_needs_upcast in previous commit even needed? if the tensor is already the desired type, the operation would be no-op.

[brkirch]

x = model.decode_first_stage(x.to(devices.dtype_vae) if devices.unet_needs_upcast else x)

Is this right?

Yes, that is what is needed. Otherwise full previews can cause an exception with --upcast-sampling.

Also I wanted to test it but never got to it, are all those if devices.unet_needs_upcast in previous commit even needed? if the tensor is already the desired type, the operation would be no-op.

The real problem is that potentially the dtype isn't necessarily going to already be correct, but casting without --upcast-sampling will have the potential to cause autocast to do an unnecessary cast. If --upcast-sampling is enabled then the casts are done regardless of autocast to ensure that the casting is done reliably, since autocast in ROCm is sometimes too unreliable to use Stable Diffusion without --upcast-sampling or --no-half.

In other words, if devices.unet_needs_upcast is to ensure that --upcast-sampling gives maximum compatiblity at the potential cost of some extra casting whereas omitting --upcast-sampling should not attempt those casts because autocast will handle it with better performance.

If we're talking about what must have the check, apply_model definitely needs the unet_needs_upcast check because that is what downcasts and upcasts in and out of the UNet. The rest of these unet_needs_upcast checks probably aren't an absolute requirement but it'd be necessary to check what the dtypes are with and without --no-half and with and without --no-half-vae when autocast is enabled to be sure that removing the checks isn't harming performance at all or changing the precision of some ops (which could change seeds).