Fix SD2 loading

[huchenlei]

Description

Closes #16031

The bug was originally introduced in #15820

The original PR trys to map

CondFunc('sgm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast)

to

def timestep_embedding_cast_result(orig_func, timesteps, *args, **kwargs):
    if devices.unet_needs_upcast and timesteps.dtype == torch.int64:
        dtype = torch.float32
    else:
        dtype = devices.dtype_unet
    return orig_func(timesteps, *args, **kwargs).to(dtype=dtype)
CondFunc('sgm.modules.diffusionmodules.openaimodel.timestep_embedding', timestep_embedding_cast_result)

However, if unet_needs_upcast is False, we should directly call the original function by the original impl. This PR patches this logic mismatch.

Checklist:

• I have read contributing wiki page
• I have performed a self-review of my own code
• My code follows the style guidelines
• My code passes tests

[huchenlei]

Closing as this is not a proper fix.