search

MStypulkowski / diffused-heads

Official repository for Diffused Heads: Diffusion Models Beat GANs on Talking-Face Generation
Other
463 stars 31 forks source link

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu) #7

Open hoalarious opened 1 year ago

hoalarious commented 1 year ago

Getting an error:

Traceback (most recent call last):
  File "/kaggle/working/diffused-heads-colab/sample.py", line 24, in main
    samples = diffusion.sample(id_frame, audio_emb.unsqueeze(0), **args.unet)
  File "/kaggle/working/diffused-heads-colab/diffusion.py", line 42, in sample
    sample_frame = self.sample_loop(xT[:, i].to(x_cond.device), x_cond, motion_frames, audio_emb[:, audio_ids])
  File "/kaggle/working/diffused-heads-colab/diffusion.py", line 54, in sample_loop
    mean, logvar = self.get_p_params(xt, timesteps_ids, nn_out)
  File "/kaggle/working/diffused-heads-colab/diffusion.py", line 67, in get_p_params
    p_logvar = nu * self.expand(torch.log(self.beta[timesteps])) + (1 - nu) * self.expand(self.log_beta_tilde_clipped[timesteps])
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Trying to get this working on colab/kaggle.

Need to setup my local environment to debug this better. Wondering if anyone has a fix for this in the mean time?

Bebaam commented 1 year ago

Not sure how it impacts performance, but at least for testing you can just put self.beta[timesteps] and self.log_beta_tilde_clipped to your GPU like this in the line before: self.beta = self.beta.to(xt.device) and self.log_beta_tilde_clipped = self.log_beta_tilde_clipped.to(xt.device)

The same you have to do for the subsequent errors, then you should be able to successfully sample.

hoalarious commented 1 year ago

Thanks I was hoping for a more elegant solution but probably not a good idea to get tied down on it at this stage. This talking head implementation is rather computational intensive. Taking 16 minutes for a 2 second clip. Going to test it on longer outputs to see if it gets better. I don't think that's from the fixes you suggested.

Got a working colab demo if anyone wants to try it: https://github.com/hoalarious/diffused-heads-colab/blob/main/diffused_heads_colab.ipynb