Enable ONNX export of GPU loaded SVD/SVD-XT UNet models

* Unpack num_frames scalar if created as a (CPU) tensor in forward path
  Avoids mixed use of CPU and CUDA tensors which is unsupported by torch.nn ops

Signed-off-by: Rajeev Rao <[email protected]>

Author: rajeevsrao

src/diffusers/models/unet_spatio_temporal_condition.py

@@ -397,6 +397,8 @@ def forward(
 
         # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         batch_size, num_frames = sample.shape[:2]
+        if torch.is_tensor(num_frames):
+            num_frames = num_frames.item()
         timesteps = timesteps.expand(batch_size)
 
         t_emb = self.time_proj(timesteps)