turboedit commited on
Commit
c813f05
1 Parent(s): 3347638

Upload utils.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. utils.py +5 -4
utils.py CHANGED
@@ -485,15 +485,16 @@ def find_index(timesteps, timestep):
485
  return i
486
  return -1
487
 
 
488
  map_timpstep_to_index = {
489
  torch.tensor(799): 0,
490
  torch.tensor(599): 1,
491
  torch.tensor(399): 2,
492
  torch.tensor(199): 3,
493
- torch.tensor(799, device='cuda:0'): 0,
494
- torch.tensor(599, device='cuda:0'): 1,
495
- torch.tensor(399, device='cuda:0'): 2,
496
- torch.tensor(199, device='cuda:0'): 3,
497
  }
498
 
499
  def step_save_latents(
 
485
  return i
486
  return -1
487
 
488
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
489
  map_timpstep_to_index = {
490
  torch.tensor(799): 0,
491
  torch.tensor(599): 1,
492
  torch.tensor(399): 2,
493
  torch.tensor(199): 3,
494
+ torch.tensor(799, device=device): 0,
495
+ torch.tensor(599, device=device): 1,
496
+ torch.tensor(399, device=device): 2,
497
+ torch.tensor(199, device=device): 3,
498
  }
499
 
500
  def step_save_latents(